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 f932a6d..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,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 diff --git a/common/utils/console_client.py b/common/utils/console_client.py index 638426b..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,6 +57,8 @@ def reserve_quota( organization_slug: str, feature: str | list[str], amount: int = 1, + scope_type: str = FeatureUsageScope.ORGANIZATION, + scope_id: str | None = None, ) -> tuple[bool, str | None]: """Reserve quota for one or more features. @@ -64,14 +67,19 @@ 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, + "scope_type": scope_type, + } + if scope_id is not None: + payload["scope_id"] = str(scope_id) + try: response = requests.post( endpoint, - json={ - "organization": organization_slug, - "feature": feature, - "amount": amount, - }, + json=payload, headers=self._headers(), timeout=self.timeout, ) @@ -106,17 +114,24 @@ def release_quota( organization_slug: str, feature: str | list[str], amount: int = 1, + 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" + payload = { + "organization": organization_slug, + "feature": feature, + "amount": amount, + "scope_type": scope_type, + } + if scope_id is not None: + payload["scope_id"] = str(scope_id) + try: requests.post( endpoint, - json={ - "organization": organization_slug, - "feature": feature, - "amount": amount, - }, + json=payload, headers=self._headers(), timeout=self.timeout, ) @@ -132,19 +147,26 @@ def get_quota( self, organization_slug: str, feature: str | list[str], + scope_type: str = FeatureUsageScope.ORGANIZATION, + scope_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/quota" + payload = { + "organization": organization_slug, + "feature": feature, + "scope_type": scope_type, + } + if scope_id is not None: + payload["scope_id"] = str(scope_id) + try: response = requests.post( endpoint, - json={ - "organization": organization_slug, - "feature": feature, - }, + json=payload, headers=self._headers(), timeout=self.timeout, )