fix: user scoped feature usage quota

common/apps/billing/constants.py

@@ -0,0 +1,11 @@
+class FeatureUsageScope:
+    ORGANIZATION = "organization"
+    USER = "user"
+    SPACE = "space"
+
+
+class FeatureCode:
+    SPACE_MAX_COUNT = "space.max_count"
+    DEVICE_MAX_COUNT = "device.max_count"
+    DASHBOARD_MAX_COUNT = "dashboard.max_count"
+    DASHBOARD_UPDATE_COUNT = "dashboard.update_count"

common/apps/billing/mixins.py

@@ -2,22 +2,24 @@
 
 Usage::
 
+    from common.apps.billing.constants import FeatureCode
+
     class DeviceQuota(BaseQuota):
         rules = {
-            "create": "device.max_count",
-            "bulk_create": "device.max_count",
-            "list": "device.read_limit",
-            "retrieve": "device.read_limit",
+            "create": FeatureCode.DEVICE_MAX_COUNT,
+            "bulk_create": FeatureCode.DEVICE_MAX_COUNT,
         }
 
     class DeviceViewSet(QuotaMixin, viewsets.ModelViewSet):
         quota_classes = [DeviceQuota]
 
-Rules can use a single feature code or a list of feature codes.
+Rules can use a single feature code, a list of feature codes, or scoped rule
+objects like ``{"feature": FeatureCode.DEVICE_MAX_COUNT, "scope": scope}``.
 """
 
 import logging
 
+from common.apps.billing.constants import FeatureUsageScope
 from rest_framework.exceptions import PermissionDenied
 
 from common.utils.console_client import console_client
@@ -27,6 +29,8 @@ class DeviceViewSet(QuotaMixin, viewsets.ModelViewSet):
 
 class BaseQuota:
     rules = {}
+    scope = FeatureUsageScope.USER
+    feature_scopes = {}
     reserve_actions = {"create", "bulk_create"}
     release_actions = set()
     method_action_map = {
@@ -48,19 +52,140 @@ def get_rule(self, request, view):
         action = self.get_action(request, view)
         if not action:
             return None
-        return self.rules.get(action)
+        rule = self.rules.get(action)
+        if rule is not None:
+            return rule
 
-    def get_features(self, request, view):
+        for actions, rule in self.rules.items():
+            if isinstance(actions, (tuple, frozenset)) and action in actions:
+                return rule
+
+        return None
+
+    def _normalize_rule_item(self, item):
+        if isinstance(item, str):
+            return {
+                "feature": item,
+                "scope": self.feature_scopes.get(item, self.scope),
+            }
+
+        if isinstance(item, dict):
+            if "feature" in item:
+                feature = item["feature"]
+                return {
+                    "feature": feature,
+                    "scope": item.get(
+                        "scope",
+                        self.feature_scopes.get(feature, self.scope),
+                    ),
+                }
+
+            return [
+                {
+                    "feature": feature,
+                    "scope": scope,
+                }
+                for feature, scope in item.items()
+            ]
+
+        if isinstance(item, (list, tuple)) and len(item) == 2:
+            return {
+                "feature": item[0],
+                "scope": item[1],
+            }
+
+        return None
+
+    def get_feature_rules(self, request, view):
         rule = self.get_rule(request, view)
         if not rule:
             return []
-        if isinstance(rule, str):
-            return [rule]
-        return list(dict.fromkeys(rule))
+
+        items = rule if isinstance(rule, list) else [rule]
+        feature_rules = []
+
+        for item in items:
+            normalized = self._normalize_rule_item(item)
+            if normalized is None:
+                continue
+            if isinstance(normalized, list):
+                feature_rules.extend(normalized)
+            else:
+                feature_rules.append(normalized)
+
+        deduped_rules = []
+        seen = set()
+        for feature_rule in feature_rules:
+            key = (feature_rule["feature"], feature_rule["scope"])
+            if key in seen:
+                continue
+            seen.add(key)
+            deduped_rules.append(feature_rule)
+
+        return deduped_rules
+
+    def get_features(self, request, view):
+        feature_rules = self.get_feature_rules(request, view)
+        if not feature_rules:
+            return []
+
+        features = [feature_rule["feature"] for feature_rule in feature_rules]
+        return list(dict.fromkeys(features))
 
     def get_org_slug(self, request, view):
         tenant = getattr(request, "tenant", None)
-        return getattr(tenant, "slug_name", None)
+        return getattr(tenant, "slug_name", None) or request.headers.get(
+            "X-Organization"
+        )
+
+    def get_user_id(self, request, view):
+        user_id = request.headers.get("X-User-ID")
+        if user_id:
+            return user_id
+
+        user = getattr(request, "user", None)
+        if getattr(user, "is_authenticated", False):
+            return getattr(user, "id", None)
+        return None
+
+    def get_scope_id(self, request, view, scope_type):
+        if scope_type == FeatureUsageScope.ORGANIZATION:
+            return None
+        if scope_type == FeatureUsageScope.USER:
+            return self.get_user_id(request, view)
+        if scope_type == FeatureUsageScope.SPACE:
+            space_slug = request.headers.get("X-Space")
+            if not space_slug:
+                return None
+
+            from common.apps.space.models import Space
+
+            return (
+                Space.objects.filter(slug_name=space_slug)
+                .values_list("id", flat=True)
+                .first()
+            )
+        if hasattr(view, "get_quota_scope_id"):
+            return view.get_quota_scope_id(request, self, scope_type)
+        return None
+
+    def get_feature_scope(self, request, view, feature):
+        return self.feature_scopes.get(feature, self.scope)
+
+    def _group_features_by_scope(self, request, view, feature_rules):
+        grouped_features = {}
+
+        for feature_rule in feature_rules:
+            feature = feature_rule["feature"]
+            scope_type = feature_rule["scope"]
+            scope_id = self.get_scope_id(request, view, scope_type)
+            if scope_type != FeatureUsageScope.ORGANIZATION and not scope_id:
+                self.message = f"Scope id is required for quota scope '{scope_type}'."
+                return None
+
+            grouped_features.setdefault((scope_type, scope_id), []).append(feature)
+
+        return grouped_features
 
     def get_amount(self, request, view):
         if hasattr(view, "get_quota_amount"):
@@ -99,19 +224,37 @@ def release(self, request, view, amount=None):
             key = (
                 id(reservation["quota"]),
                 reservation["org_slug"],
+                reservation["scope_type"],
+                reservation["scope_id"],
                 reservation["feature"],
                 reservation["amount"],
             )
             if key in released:
                 continue
 
-            group_key = (reservation["org_slug"], reservation["amount"])
+            group_key = (
+                reservation["org_slug"],
+                reservation["scope_type"],
+                reservation["scope_id"],
+                reservation["amount"],
+            )
             grouped_features.setdefault(group_key, []).append(reservation["feature"])
             released.add(key)
 
-        for (org_slug, reserved_amount), features in grouped_features.items():
+        for (
+            org_slug,
+            scope_type,
+            scope_id,
+            reserved_amount,
+        ), features in grouped_features.items():
             try:
-                console_client.release_quota(org_slug, features, reserved_amount)
+                console_client.release_quota(
+                    org_slug,
+                    features,
+                    reserved_amount,
+                    scope_type=scope_type,
+                    scope_id=scope_id,
+                )
             except Exception:  # noqa: BLE001
                 logger.warning("release_quota failed for %s/%s", org_slug, features)
 
@@ -127,46 +270,74 @@ def check(self, request, view):
         return self.has_quota(request, view)
 
     def _check_features(self, request, view, amount):
-        features = self.get_features(request, view)
-        if not features:
+        feature_rules = self.get_feature_rules(request, view)
+        if not feature_rules:
             return True
 
         org_slug = self.get_org_slug(request, view)
         if not org_slug:
             return True
 
-        reserved, error = console_client.reserve_quota(org_slug, features, amount)
-        if not reserved:
-            self.message = error or self.message
-            self.release(request, view)
+        grouped_features = self._group_features_by_scope(request, view, feature_rules)
+        if grouped_features is None:
             return False
 
-        for feature in features:
-            if amount > 0:
-                view._quota_reserved.append(
-                    {
-                        "quota": self,
-                        "org_slug": org_slug,
-                        "feature": feature,
-                        "amount": amount,
-                    }
-                )
+        for (scope_type, scope_id), scoped_features in grouped_features.items():
+            reserved, error = console_client.reserve_quota(
+                org_slug,
+                scoped_features,
+                amount,
+                scope_type=scope_type,
+                scope_id=scope_id,
+            )
+            if not reserved:
+                self.message = error or self.message
+                self.release(request, view)
+                return False
+
+            for feature in scoped_features:
+                if amount > 0:
+                    view._quota_reserved.append(
+                        {
+                            "quota": self,
+                            "org_slug": org_slug,
+                            "scope_type": scope_type,
+                            "scope_id": scope_id,
+                            "feature": feature,
+                            "amount": amount,
+                        }
+                    )
 
         return True
 
     def _release_features(self, request, view, amount):
-        features = self.get_features(request, view)
-        if not features:
+        feature_rules = self.get_feature_rules(request, view)
+        if not feature_rules:
             return None
 
         org_slug = self.get_org_slug(request, view)
         if not org_slug:
             return None
 
-        try:
-            console_client.release_quota(org_slug, features, amount)
-        except Exception:  # noqa: BLE001
-            logger.warning("release_quota failed for %s/%s", org_slug, features)
+        grouped_features = self._group_features_by_scope(request, view, feature_rules)
+        if grouped_features is None:
+            return None
+
+        for (scope_type, scope_id), scoped_features in grouped_features.items():
+            try:
+                console_client.release_quota(
+                    org_slug,
+                    scoped_features,
+                    amount,
+                    scope_type=scope_type,
+                    scope_id=scope_id,
+                )
+            except Exception:  # noqa: BLE001
+                logger.warning(
+                    "release_quota failed for %s/%s",
+                    org_slug,
+                    scoped_features,
+                )
 
         return None