From 14f5f45fae7d860d784c843653957495538d3ac5 Mon Sep 17 00:00:00 2001 From: ngovinh2k2 Date: Fri, 17 Jul 2026 09:58:34 +0700 Subject: [PATCH 1/2] fix: user scoped feature usage quota --- common/apps/billing/mixins.py | 205 ++++++++++++++++++++++++++++----- common/utils/console_client.py | 43 ++++--- 2 files changed, 203 insertions(+), 45 deletions(-) diff --git a/common/apps/billing/mixins.py b/common/apps/billing/mixins.py index f932a6d..25a267d 100644 --- a/common/apps/billing/mixins.py +++ b/common/apps/billing/mixins.py @@ -27,6 +27,8 @@ class DeviceViewSet(QuotaMixin, viewsets.ModelViewSet): class BaseQuota: rules = {} + scope = "user" + feature_scopes = {} reserve_actions = {"create", "bulk_create"} release_actions = set() method_action_map = { @@ -48,19 +50,125 @@ 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_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 = feature_rule["scope"] + if scope == "organization": + user_id = None + elif scope == "user": + user_id = self.get_user_id(request, view) + if not user_id: + self.message = "User is required for quota." + return None + else: + self.message = f"Invalid quota scope '{scope}'." + return None + + grouped_features.setdefault(user_id, []).append(feature) + + return grouped_features def get_amount(self, request, view): if hasattr(view, "get_quota_amount"): @@ -99,19 +207,29 @@ def release(self, request, view, amount=None): key = ( id(reservation["quota"]), reservation["org_slug"], + reservation["user_id"], reservation["feature"], reservation["amount"], ) if key in released: continue - group_key = (reservation["org_slug"], reservation["amount"]) + group_key = ( + reservation["org_slug"], + reservation["user_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, user_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, + user_id=user_id, + ) except Exception: # noqa: BLE001 logger.warning("release_quota failed for %s/%s", org_slug, features) @@ -127,46 +245,71 @@ 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 user_id, scoped_features in grouped_features.items(): + reserved, error = console_client.reserve_quota( + org_slug, + scoped_features, + amount, + user_id=user_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, + "user_id": user_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 user_id, scoped_features in grouped_features.items(): + try: + console_client.release_quota( + org_slug, + scoped_features, + amount, + user_id=user_id, + ) + except Exception: # noqa: BLE001 + logger.warning( + "release_quota failed for %s/%s", + org_slug, + scoped_features, + ) return None diff --git a/common/utils/console_client.py b/common/utils/console_client.py index 4068819..c1d76e4 100644 --- a/common/utils/console_client.py +++ b/common/utils/console_client.py @@ -56,6 +56,7 @@ def reserve_quota( organization_slug: str, feature: str | list[str], amount: int = 1, + user_id: str | None = None, ) -> tuple[bool, str | None]: """Reserve quota for one or more features. @@ -64,14 +65,18 @@ def reserve_quota( billing infra problems never block the core product. """ endpoint = f"{self.base_url}/billing/internal/quota/reserve" + payload = { + "organization": organization_slug, + "feature": feature, + "amount": amount, + } + if user_id is not None: + payload["user_id"] = str(user_id) + try: response = requests.post( endpoint, - json={ - "organization": organization_slug, - "feature": feature, - "amount": amount, - }, + json=payload, headers=self._headers(), timeout=self.timeout, ) @@ -106,17 +111,22 @@ def release_quota( organization_slug: str, feature: str | list[str], amount: int = 1, + user_id: str | None = None, ) -> None: """Release previously reserved quota for one or more features.""" endpoint = f"{self.base_url}/billing/internal/quota/release" + payload = { + "organization": organization_slug, + "feature": feature, + "amount": amount, + } + if user_id is not None: + payload["user_id"] = str(user_id) + try: requests.post( endpoint, - json={ - "organization": organization_slug, - "feature": feature, - "amount": amount, - }, + json=payload, headers=self._headers(), timeout=self.timeout, ) @@ -132,19 +142,24 @@ def get_quota( self, organization_slug: str, feature: str | list[str], + user_id: str | None = None, ) -> tuple[dict | None, str | None]: """View current quota usage for one or more features. Returns ``(data, error)`` matching the reserve/release pattern. """ endpoint = f"{self.base_url}/billing/internal/quota/view" + payload = { + "organization": organization_slug, + "feature": feature, + } + if user_id is not None: + payload["user_id"] = str(user_id) + try: response = requests.post( endpoint, - json={ - "organization": organization_slug, - "feature": feature, - }, + json=payload, headers=self._headers(), timeout=self.timeout, ) From f174fbfe6292bbce557fe8d254deb20078119cf2 Mon Sep 17 00:00:00 2001 From: ngovinh2k2 Date: Fri, 17 Jul 2026 11:02:53 +0700 Subject: [PATCH 2/2] fix: update logic code --- common/apps/billing/constants.py | 11 +++++ common/apps/billing/mixins.py | 80 +++++++++++++++++++++----------- common/utils/console_client.py | 25 ++++++---- 3 files changed, 81 insertions(+), 35 deletions(-) create mode 100644 common/apps/billing/constants.py diff --git a/common/apps/billing/constants.py b/common/apps/billing/constants.py new file mode 100644 index 0000000..7f5da0a --- /dev/null +++ b/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" diff --git a/common/apps/billing/mixins.py b/common/apps/billing/mixins.py index 25a267d..10882fe 100644 --- a/common/apps/billing/mixins.py +++ b/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,7 +29,7 @@ class DeviceViewSet(QuotaMixin, viewsets.ModelViewSet): class BaseQuota: rules = {} - scope = "user" + scope = FeatureUsageScope.USER feature_scopes = {} reserve_actions = {"create", "bulk_create"} release_actions = set() @@ -146,6 +148,27 @@ def get_user_id(self, request, view): 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) @@ -154,19 +177,13 @@ def _group_features_by_scope(self, request, view, feature_rules): for feature_rule in feature_rules: feature = feature_rule["feature"] - scope = feature_rule["scope"] - if scope == "organization": - user_id = None - elif scope == "user": - user_id = self.get_user_id(request, view) - if not user_id: - self.message = "User is required for quota." - return None - else: - self.message = f"Invalid quota scope '{scope}'." + 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(user_id, []).append(feature) + grouped_features.setdefault((scope_type, scope_id), []).append(feature) return grouped_features @@ -207,7 +224,8 @@ def release(self, request, view, amount=None): key = ( id(reservation["quota"]), reservation["org_slug"], - reservation["user_id"], + reservation["scope_type"], + reservation["scope_id"], reservation["feature"], reservation["amount"], ) @@ -216,19 +234,26 @@ def release(self, request, view, amount=None): group_key = ( reservation["org_slug"], - reservation["user_id"], + reservation["scope_type"], + reservation["scope_id"], reservation["amount"], ) grouped_features.setdefault(group_key, []).append(reservation["feature"]) released.add(key) - for (org_slug, user_id, 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, - user_id=user_id, + scope_type=scope_type, + scope_id=scope_id, ) except Exception: # noqa: BLE001 logger.warning("release_quota failed for %s/%s", org_slug, features) @@ -257,12 +282,13 @@ def _check_features(self, request, view, amount): if grouped_features is None: return False - for user_id, scoped_features in grouped_features.items(): + for (scope_type, scope_id), scoped_features in grouped_features.items(): reserved, error = console_client.reserve_quota( org_slug, scoped_features, amount, - user_id=user_id, + scope_type=scope_type, + scope_id=scope_id, ) if not reserved: self.message = error or self.message @@ -275,7 +301,8 @@ def _check_features(self, request, view, amount): { "quota": self, "org_slug": org_slug, - "user_id": user_id, + "scope_type": scope_type, + "scope_id": scope_id, "feature": feature, "amount": amount, } @@ -296,13 +323,14 @@ def _release_features(self, request, view, amount): if grouped_features is None: return None - for user_id, scoped_features in grouped_features.items(): + for (scope_type, scope_id), scoped_features in grouped_features.items(): try: console_client.release_quota( org_slug, scoped_features, amount, - user_id=user_id, + scope_type=scope_type, + scope_id=scope_id, ) except Exception: # noqa: BLE001 logger.warning( diff --git a/common/utils/console_client.py b/common/utils/console_client.py index 7d39645..4004e87 100644 --- a/common/utils/console_client.py +++ b/common/utils/console_client.py @@ -1,6 +1,7 @@ import logging import requests +from common.apps.billing.constants import FeatureUsageScope from requests.exceptions import RequestException logger = logging.getLogger(__name__) @@ -56,7 +57,8 @@ def reserve_quota( organization_slug: str, feature: str | list[str], amount: int = 1, - user_id: str | None = None, + scope_type: str = FeatureUsageScope.ORGANIZATION, + scope_id: str | None = None, ) -> tuple[bool, str | None]: """Reserve quota for one or more features. @@ -69,9 +71,10 @@ def reserve_quota( "organization": organization_slug, "feature": feature, "amount": amount, + "scope_type": scope_type, } - if user_id is not None: - payload["user_id"] = str(user_id) + if scope_id is not None: + payload["scope_id"] = str(scope_id) try: response = requests.post( @@ -111,7 +114,8 @@ def release_quota( organization_slug: str, feature: str | list[str], amount: int = 1, - user_id: str | None = None, + scope_type: str = FeatureUsageScope.ORGANIZATION, + scope_id: str | None = None, ) -> None: """Release previously reserved quota for one or more features.""" endpoint = f"{self.base_url}/billing/internal/quota/release" @@ -119,9 +123,10 @@ def release_quota( "organization": organization_slug, "feature": feature, "amount": amount, + "scope_type": scope_type, } - if user_id is not None: - payload["user_id"] = str(user_id) + if scope_id is not None: + payload["scope_id"] = str(scope_id) try: requests.post( @@ -142,7 +147,8 @@ def get_quota( self, organization_slug: str, feature: str | list[str], - user_id: str | None = None, + scope_type: str = FeatureUsageScope.ORGANIZATION, + scope_id: str | None = None, ) -> tuple[dict | None, str | None]: """View current quota usage for one or more features. @@ -152,9 +158,10 @@ def get_quota( payload = { "organization": organization_slug, "feature": feature, + "scope_type": scope_type, } - if user_id is not None: - payload["user_id"] = str(user_id) + if scope_id is not None: + payload["scope_id"] = str(scope_id) try: response = requests.post(