diff --git a/rhodecode/apps/admin/tests/test_admin_repos.py b/rhodecode/apps/admin/tests/test_admin_repos.py index 99f75b2c..b2763f96 100644 --- a/rhodecode/apps/admin/tests/test_admin_repos.py +++ b/rhodecode/apps/admin/tests/test_admin_repos.py @@ -106,6 +106,11 @@ class TestAdminRepos(object): monkeypatch.setitem(app.CONFIG, "user.quotas.enabled", True) + user: User = User.get_by_username(TEST_USER_REGULAR_LOGIN) + existing_repo = backend.create_repo() + existing_repo.user = user + Session().commit() + session = login_user_session(self.app, TEST_USER_REGULAR_LOGIN, TEST_USER_REGULAR_PASS) csrf_token = auth.get_csrf_token(session) @@ -114,9 +119,8 @@ class TestAdminRepos(object): user: User = User.get_by_username(TEST_USER_REGULAR_LOGIN) user.update_quotas( - max_disk_space_mb_allowance=200, - max_repository_count_allowance=2, - current_repos_count=2, + max_disk_space_mb_allowance=100, + max_repository_count_allowance=1, used_disk_space_bytes=0, ) @@ -153,9 +157,8 @@ class TestAdminRepos(object): user: User = User.get_by_username(TEST_USER_REGULAR_LOGIN) user.update_quotas( - max_disk_space_mb_allowance=200, - max_repository_count_allowance=2, - current_repos_count=0, + max_disk_space_mb_allowance=100, + max_repository_count_allowance=1, used_disk_space_bytes=0, ) @@ -172,7 +175,8 @@ class TestAdminRepos(object): self.assert_repository_is_created_correctly(repo_name, "quota", backend) user: User = User.get_by_username(TEST_USER_REGULAR_LOGIN) - assert user.get_quotas()["current_repos_count"] == 1 + user_repos = Repository.get_all_repos(user.user_id) + assert len(user_repos) == 1 def test_create_numeric_name(self, autologin_user, backend, csrf_token): numeric_repo = "1234" diff --git a/rhodecode/apps/repository/tests/test_repo_settings_advanced.py b/rhodecode/apps/repository/tests/test_repo_settings_advanced.py index de404afb..5a8690e1 100644 --- a/rhodecode/apps/repository/tests/test_repo_settings_advanced.py +++ b/rhodecode/apps/repository/tests/test_repo_settings_advanced.py @@ -127,9 +127,8 @@ class TestAdminRepoSettingsAdvanced(object): user: User = User.get_by_username(TEST_USER_REGULAR_LOGIN) user.update_quotas( - max_disk_space_mb_allowance=200, - max_repository_count_allowance=2, - current_repos_count=1, + max_disk_space_mb_allowance=100, + max_repository_count_allowance=1, used_disk_space_bytes=0, ) @@ -151,7 +150,8 @@ class TestAdminRepoSettingsAdvanced(object): # check if repo was deleted from db assert RepoModel().get_by_repo_name(repo_name) is None assert not repo_on_filesystem(repo_name_str) - assert user.get_quotas()["current_repos_count"] == 0 + user_repos = Repository.get_all_repos(user.user_id) + assert len(user_repos) == 0 @pytest.mark.parametrize("suffix", ["", "ąęł", "123"], ids=no_newline_id_generator) def test_delete_broken_repo(self, autologin_user, backend, suffix, csrf_token): diff --git a/rhodecode/apps/repository/views/repo_settings_advanced.py b/rhodecode/apps/repository/views/repo_settings_advanced.py index feff3b09..6b81ca3a 100644 --- a/rhodecode/apps/repository/views/repo_settings_advanced.py +++ b/rhodecode/apps/repository/views/repo_settings_advanced.py @@ -158,9 +158,6 @@ class RepoSettingsAdvancedView(RepoAppView): self.request.POST.get("backup", "false") ) # if this option is absent, this means that on FE side checkbox was unchecked delete_res = RepoModel().delete(self.db_repo, fs_backup=backup) - repo_owner = self.db_repo.user - quota_model = UserQuotaModel(repo_owner) - quota_model.decrement_repo_count() fork_action = delete_res.fork_action_result.action fork_cnt = delete_res.fork_action_result.fork_count diff --git a/rhodecode/lib/celerylib/tasks.py b/rhodecode/lib/celerylib/tasks.py index 4ab112a4..83a36d5e 100644 --- a/rhodecode/lib/celerylib/tasks.py +++ b/rhodecode/lib/celerylib/tasks.py @@ -280,9 +280,6 @@ def create_repo(form_data, cur_user): repo=audit_logger.RepoWrap(repo_name=repo_name, repo_id=repo_id), ) - quota_model = UserQuotaModel(cur_user) - quota_model.increment_repo_count() - Session().commit() PermissionModel().trigger_permission_flush() diff --git a/rhodecode/model/quota_model.py b/rhodecode/model/quota_model.py index f9c13916..2fa23bf6 100644 --- a/rhodecode/model/quota_model.py +++ b/rhodecode/model/quota_model.py @@ -1,40 +1,22 @@ import logging -from functools import wraps -from rhodecode.model.db import User +from sqlalchemy import func, select + import rhodecode as app - +from rhodecode.model.db import User, Repository +from rhodecode.model.meta import Session log = logging.getLogger(__name__) -def always_valid_for_super_admin(f): - @wraps(f) - def inner(self, *args, **kwargs): - if self.user.is_admin: - return True - return f(self, *args, **kwargs) - - return inner - - -def skip_for_super_admin(f): - @wraps(f) - def inner(self, *args, **kwargs): - if self.user.is_admin: - log.debug(f"Skipping {f.__name__} for super admin.") - return - return f(self, *args, **kwargs) - - return inner - - class UserQuotaModel: """ NOTE: quotas are applied only when a user registers itself or the admin adds a user, in plugins there are more ways to add a user. """ + UNLIMITED = float("inf") + class QuotaError(Exception): pass @@ -53,52 +35,44 @@ class UserQuotaModel: self.user.update_quotas( max_disk_space_mb_allowance=self.default_max_disk_space_mb_allowance, max_repository_count_allowance=self.default_max_repository_count_allowance, - current_repos_count=0, used_disk_space_bytes=0, ) - @always_valid_for_super_admin + def _init_quotas_for_old_users(self): + if not self.enabled: + log.debug("quotas are disabled") + return + + log.debug("Initializing quotas for old users.") + + self.user.update_quotas( + max_disk_space_mb_allowance=self.UNLIMITED, + max_repository_count_allowance=self.UNLIMITED, + used_disk_space_bytes=0, + ) + + def user_repositories_count(self): + return Session().scalar( + select(func.count()).select_from(Repository).where(Repository.user_id == self.user.user_id) + ) + def is_repo_creation_allowed(self): if not self.enabled: return True if self._user_quotas_valid(): u_quotas = self.user.get_quotas() - return u_quotas.get("max_repository_count_allowance") > u_quotas.get("current_repos_count") + return u_quotas.get("max_repository_count_allowance") > self.user_repositories_count() - log.warning("User quotas not present, probably old user, allowing repo creation.") - return True - - @skip_for_super_admin - def increment_repo_count(self): - if not self.enabled: - return - - if self._user_quotas_valid(): - u_quotas = self.user.get_quotas() - new_count = u_quotas["current_repos_count"] + 1 - self.user.update_quotas(current_repos_count=new_count) - - log.warning("User quotas not present, probably old user, skipping.") - - @skip_for_super_admin - def decrement_repo_count(self): - if not self.enabled: - return - - u_quotas = self.user.get_quotas() - if self._user_quotas_valid() and u_quotas["current_repos_count"] > 0: - new_count = u_quotas["current_repos_count"] - 1 - self.user.update_quotas(current_repos_count=new_count) - - log.warning("User quotas not present, probably old user, skipping.") + log.warning("User quotas not present, probably old user.") + self._init_quotas_for_old_users() + 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", - "current_repos_count", "used_disk_space_bytes", ]: if u_quotas.get(key) is None: