feature: initial implementation of qoutas for a user, repo count part

rhodecode/model/db.py

@@ -1256,6 +1256,15 @@ class User(Base, BaseModel):
         if user:
             return user
 
+    def update_quotas(self, **kwargs):
+        quotas = self.user_data.get("quotas", {})
+        quotas.update(**kwargs)
+        self.update_userdata(quotas=quotas)
+        Session().commit()
+
+    def get_quotas(self):
+        return self.user_data.get("quotas", {})
+
     def update_userdata(self, **kwargs):
         usr = self
         old = usr.user_data

rhodecode/model/quota_model.py

@@ -0,0 +1,110 @@
+import logging
+from functools import wraps
+
+from rhodecode.model.db import User
+import rhodecode as app
+
+
+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.
+    """
+
+    class QuotaError(Exception):
+        pass
+
+    def __init__(self, user: User):
+        self._settings = app.CONFIG
+        self.enabled = self._settings.get("user.quotas.enabled")
+        self.default_max_disk_space_mb_allowance = self._settings.get("user.quotas.disk_space_mb")
+        self.default_max_repository_count_allowance = self._settings.get("user.quotas.repository_count")
+        self.user = user
+
+    def init_quotas(self):
+        if not self.enabled:
+            log.debug("quotas are disabled")
+            return
+
+        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 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")
+
+        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.")
+
+    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:
+                log.error(
+                    f"User {self.user.username} has no quota value for {key}, probably quota initialization was not done."
+                )
+                return False
+
+        return True

rhodecode/model/user.py

@@ -62,6 +62,7 @@ from rhodecode.model.db import (
 )
 from rhodecode.model.meta import Session
 from rhodecode.model.auth_token import AuthTokenModel
+from rhodecode.model.quota_model import UserQuotaModel
 from rhodecode.model.repo_group import RepoGroupModel
 
 log = logging.getLogger(__name__)
@@ -409,6 +410,10 @@ class UserModel(BaseModel):
                 new_user.update_userdata(language=language)
             new_user.update_userdata(notification_status=True)
 
+            if new_user.is_new_user:
+                qm = UserQuotaModel(new_user)
+                qm.init_quotas()
+
             self.sa.add(new_user)
 
             if not edit and create_repo_group:
@@ -443,6 +448,8 @@ class UserModel(BaseModel):
             form_data["extern_type"] = extern_type
             new_user = self.create(form_data)
 
+            qm = UserQuotaModel(new_user)
+            qm.init_quotas()
             self.sa.add(new_user)
             self.sa.flush()