100 lines
3.3 KiB
Python
100 lines
3.3 KiB
Python
import logging
|
|
|
|
from sqlalchemy import func, select
|
|
|
|
from rhodecode.model.db import Repository
|
|
from rhodecode.model.meta import Session
|
|
from rhodecode.model.settings import SettingsModel
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
class UserQuotaModel:
|
|
UNLIMITED = -1
|
|
DEFAULT_REPO_QUOTA = 5
|
|
DEFAULT_QUOTAS_ENABLED = False
|
|
|
|
class QuotaError(Exception):
|
|
pass
|
|
|
|
def __init__(self, user):
|
|
self._settings = SettingsModel().get_all_settings()
|
|
self.enabled = self._settings.get("rhodecode_user_quotas_enabled", UserQuotaModel.DEFAULT_QUOTAS_ENABLED)
|
|
self.max_disk_space_mb_allowance = 100 # not implemented yet
|
|
self.max_repository_count_allowance = self._settings.get(
|
|
"rhodecode_global_user_repo_quotas", UserQuotaModel.DEFAULT_REPO_QUOTA
|
|
)
|
|
self.user = user
|
|
|
|
def init_quotas(self):
|
|
from rhodecode.model.db import User # resolve circular dependency
|
|
|
|
if not self.enabled:
|
|
log.debug("quotas are disabled")
|
|
return
|
|
|
|
mb_allowance = self.max_disk_space_mb_allowance
|
|
repo_count_allowance = self.max_repository_count_allowance
|
|
if self.user.admin or self.user.username in [
|
|
User.AI_USER,
|
|
User.DEFAULT_USER,
|
|
User.GHOST_USER,
|
|
User.SECURITY_USER,
|
|
]:
|
|
mb_allowance = self.UNLIMITED
|
|
repo_count_allowance = self.UNLIMITED
|
|
|
|
self.user.update_quotas(
|
|
max_disk_space_mb_allowance=mb_allowance,
|
|
max_repository_count_allowance=repo_count_allowance,
|
|
)
|
|
|
|
@property
|
|
def user_repositories_count(self) -> int:
|
|
return Session().scalar(
|
|
select(func.count()).select_from(Repository).where(Repository.user_id == self.user.user_id)
|
|
)
|
|
|
|
@property
|
|
def user_max_repository_count_allowance(self) -> int:
|
|
quotas = self.user.get_quotas()
|
|
return quotas.get("max_repository_count_allowance")
|
|
|
|
@user_max_repository_count_allowance.setter
|
|
def user_max_repository_count_allowance(self, new_value: int):
|
|
self.user.update_quotas(max_repository_count_allowance=new_value)
|
|
|
|
def formatted_user_max_repository_count_allowance(self) -> str:
|
|
max_repos = self.user_max_repository_count_allowance
|
|
if max_repos == self.UNLIMITED:
|
|
return "Unlimited"
|
|
return str(max_repos)
|
|
|
|
def is_repo_creation_allowed(self):
|
|
if not self.enabled:
|
|
return True
|
|
|
|
if self._user_quotas_valid():
|
|
if self.user_max_repository_count_allowance == self.UNLIMITED:
|
|
# since DB can't save float('inf'), we need to use a special case
|
|
return True
|
|
|
|
return self.user_max_repository_count_allowance > self.user_repositories_count
|
|
|
|
log.warning("User quotas not present, re-initialize quotas.")
|
|
self.init_quotas()
|
|
return self.is_repo_creation_allowed()
|
|
|
|
def _user_quotas_valid(self):
|
|
u_quotas = self.user.get_quotas()
|
|
for key in [
|
|
"max_disk_space_mb_allowance",
|
|
"max_repository_count_allowance",
|
|
]:
|
|
if u_quotas.get(key) is None:
|
|
log.error(
|
|
f"User {self.user.username} has no quota value for {key}, probably quota initialization was not done."
|
|
)
|
|
return False
|
|
|
|
return True
|