From 6d0af7a5472311f76d491ff78e1eab32dc978b57 Mon Sep 17 00:00:00 2001 From: ngovinh2k2 Date: Tue, 18 Aug 2026 09:17:30 +0700 Subject: [PATCH] feat: subscription lifecycle routing --- .../0002_seed_default_plans_and_features.py | 1 + apps/billing/services/subscription.py | 176 +++++++++++++++--- bootstrap_service/celery.py | 29 ++- 3 files changed, 175 insertions(+), 31 deletions(-) diff --git a/apps/billing/migrations/0002_seed_default_plans_and_features.py b/apps/billing/migrations/0002_seed_default_plans_and_features.py index 34bed98..1e8453e 100644 --- a/apps/billing/migrations/0002_seed_default_plans_and_features.py +++ b/apps/billing/migrations/0002_seed_default_plans_and_features.py @@ -95,6 +95,7 @@ "device.max_count": {"enabled": True, "limit_value": 10}, "space.max_count": {"enabled": True, "limit_value": 1}, "dashboard.max_count": {"enabled": True, "limit_value": 1}, + "automation.max_count": {"enabled": True, "limit_value": 0}, "dashboard.basic_widgets": {"enabled": True, "limit_value": None}, "map_view.2d": {"enabled": True, "limit_value": None}, "map_view.3d": {"enabled": True, "limit_value": None}, diff --git a/apps/billing/services/subscription.py b/apps/billing/services/subscription.py index e78f20d..adfddfd 100644 --- a/apps/billing/services/subscription.py +++ b/apps/billing/services/subscription.py @@ -1,9 +1,10 @@ import logging +from dataclasses import dataclass from datetime import timedelta from urllib.parse import urljoin -from common.apps.billing.constants import FeatureUsageScope -from common.celery.task_senders import send_task +from common.apps.billing.constants import FeatureCode, FeatureUsageScope +from common.celery.task_senders import send_subscription_task from common.utils.email_context import get_email_context, render_email_format from common.utils.send_email import send_email from django.conf import settings @@ -26,6 +27,35 @@ logger = logging.getLogger(__name__) +@dataclass(frozen=True) +class SubscriptionTask: + """A service that participates in subscription lifecycle events. + + ``task_name`` is derived from service + lifecycle (e.g. ``device_downgrade``) + so adding a service is one entry below, not two. ``required_features`` lists + the feature codes the caller must supply limits/entitlements for before the + task is dispatched (empty = no gating, e.g. console housekeeping). + """ + + service: str + required_features: tuple = () + + def task_name(self, lifecycle: str) -> str: + return f"{self.service}_{lifecycle}" + + +# Single source of truth for which services get subscription lifecycle tasks. +# To add a service: append one SubscriptionTask here. +# NOTE: "auth" gates SPACE_MAX_COUNT because the auth-service owns the Space model. +SUBSCRIPTION_SERVICES = ( + SubscriptionTask("device", (FeatureCode.DEVICE_MAX_COUNT,)), + SubscriptionTask("auth", (FeatureCode.SPACE_MAX_COUNT,)), + SubscriptionTask("dashboard", (FeatureCode.DASHBOARD_MAX_COUNT,)), + SubscriptionTask("console", ()), + SubscriptionTask("telemetry", (FeatureCode.AUTOMATION_MAX_COUNT,)), +) + + def get_free_plan_item(): return ( PlanItem.objects.select_related("plan") @@ -351,34 +381,112 @@ def get_quotas(organization, feature_codes, scope_type=None, scope_id=None): } -def _get_free_plan_limits(): - """Cached Free plan limits dict — queried once per hour.""" - cache_key = "billing:free_plan_limits" - limits = cache.get(cache_key) - if limits is not None: - return limits +def _get_plan_entitlements(plan): + if plan is None: + return {}, set() - from apps.billing.models import Plan, PlanFeature + cache_key = ":".join(("billing", "plan_entitlements", str(plan.id))) + entitlements = cache.get(cache_key) + if entitlements is not None: + return entitlements["limits"], set(entitlements["unlimited_features"]) - free_plan = Plan.objects.filter(code=PlanCodeType.FREE).first() limits = {} - if free_plan: - for pf in ( - PlanFeature.objects.filter(plan=free_plan) - .select_related("feature") - .iterator() - ): - code = pf.feature.code - if pf.limit_value is not None: - limits[code] = pf.limit_value - - cache.set(cache_key, limits, 3600) + unlimited_features = set() + for pf in ( + PlanFeature.objects.filter(plan=plan, enabled=True) + .select_related("feature") + .iterator() + ): + code = pf.feature.code + if pf.limit_value is not None: + if pf.limit_value < 0: + raise ValueError(f"plan limit {code} must be >= 0") + limits[code] = pf.limit_value + else: + unlimited_features.add(code) + + cache.set( + cache_key, + { + "limits": limits, + "unlimited_features": tuple(sorted(unlimited_features)), + }, + 3600, + ) + return limits, unlimited_features + + +def _get_plan_limits(plan): + limits, _ = _get_plan_entitlements(plan) return limits -def _send_subscription_tasks(prefix, payload): - for service in ("device", "space", "dashboard", "console", "automation"): - send_task(f"{service}_{prefix}", payload) +def _get_free_plan_limits(): + """Cached Free plan limits dict""" + from apps.billing.models import Plan + + free_plan = Plan.objects.filter(code=PlanCodeType.FREE).first() + return _get_plan_limits(free_plan) + + +def _subscription_entitlement_payload(subscription): + if ( + subscription is None + or not subscription.plan_item + or not subscription.plan_item.plan + ): + raise ValueError("subscription with plan is required for subscription upgrade") + + limits, unlimited_features = _get_plan_entitlements(subscription.plan_item.plan) + return { + "limits": limits, + "unlimited_features": sorted(unlimited_features), + } + + +def _send_subscription_tasks(lifecycle, payload): + for task in SUBSCRIPTION_SERVICES: + send_subscription_task( + task.service, lifecycle, task.task_name(lifecycle), payload + ) + + +def _required_subscription_features(): + """All feature codes any subscription task requires (across all services).""" + return {code for task in SUBSCRIPTION_SERVICES for code in task.required_features} + + +def _validate_no_negative_limits(limits, lifecycle): + negative = [ + code for code, value in limits.items() if value is not None and value < 0 + ] + if negative: + raise ValueError( + f"{lifecycle} limits must be >= 0: " + ", ".join(sorted(negative)) + ) + + +def _validate_downgrade_limits(limits): + required = _required_subscription_features() + missing = [code for code in required if code not in limits] + if missing: + raise ValueError( + "missing required downgrade limits: " + ", ".join(sorted(missing)) + ) + _validate_no_negative_limits(limits, "downgrade") + + +def _validate_upgrade_entitlements(limits, unlimited_features): + required = _required_subscription_features() + unlimited = set(unlimited_features) + missing = [ + code for code in required if code not in limits and code not in unlimited + ] + if missing: + raise ValueError( + "missing required upgrade entitlements: " + ", ".join(sorted(missing)) + ) + _validate_no_negative_limits(limits, "upgrade") def _user_display_name(user): @@ -523,6 +631,9 @@ def downgrade_to_free(organization): ) return + limits = _get_plan_limits(free_item.plan) + _validate_downgrade_limits(limits) + with transaction.atomic(): subscription = get_current_subscription(organization, for_update=True) if subscription is None: @@ -554,10 +665,10 @@ def downgrade_to_free(organization): organization.slug_name, ) - limits = _get_free_plan_limits() payload = { "org_slug": organization.slug_name, "limits": limits, + "unlimited_features": [], "downgraded_at": now.isoformat(), } @@ -607,14 +718,16 @@ def downgrade_subscription_to_free(subscription): """Downgrade a specific subscription row to Free and enforce Free limits.""" previous_plan_item = subscription.plan_item previous_period_end = subscription.period_end + limits = _get_free_plan_limits() + _validate_downgrade_limits(limits) updated = _update_subscription_to_free(subscription) if not updated: return False - limits = _get_free_plan_limits() payload = { "org_slug": subscription.organization.slug_name, "limits": limits, + "unlimited_features": [], "downgraded_at": timezone.now().isoformat(), } _send_subscription_tasks("downgrade", payload) @@ -634,7 +747,16 @@ def downgrade_subscription_to_free(subscription): def renew_subscription(organization, subscription=None, send_email_notification=False): """Enqueue upgrade tasks so services re-activate deactivated resources.""" - payload = {"org_slug": organization.slug_name} + current_subscription = subscription or get_current_subscription(organization) + entitlements = _subscription_entitlement_payload(current_subscription) + _validate_upgrade_entitlements( + entitlements["limits"], + entitlements["unlimited_features"], + ) + payload = { + "org_slug": organization.slug_name, + **entitlements, + } _send_subscription_tasks("upgrade", payload) if ( diff --git a/bootstrap_service/celery.py b/bootstrap_service/celery.py index f6e66bb..1c34588 100644 --- a/bootstrap_service/celery.py +++ b/bootstrap_service/celery.py @@ -7,6 +7,11 @@ from celery import Celery from common.celery import constants +from common.celery.routing import ( + append_unique_task_queues, + setup_subscription_task_routing, +) +from console_service.constants import CONSOLE_DOWNGRADE_TASK, CONSOLE_UPGRADE_TASK from dotenv import load_dotenv from kombu import Exchange, Queue @@ -16,27 +21,43 @@ app = Celery("bootstrap_service") app.config_from_object("django.conf:settings", namespace="CELERY") app.autodiscover_tasks() +setup_subscription_task_routing( + [ + { + "task_name": CONSOLE_DOWNGRADE_TASK, + "service": "console", + "lifecycle": "downgrade", + }, + { + "task_name": CONSOLE_UPGRADE_TASK, + "service": "console", + "lifecycle": "upgrade", + }, + ] +) + TASKS_CONSOLE = [ constants.CONSOLE_SERVICE_ADD_OR_REMOVE_SPACE, constants.CONSOLE_SERVICE_DELETE_UPLOAD_FILE, ] -existing = {queue.name: queue for queue in (app.conf.task_queues or ())} routes = dict(app.conf.task_routes or {}) +queues = [] for name in TASKS_CONSOLE: - if name not in existing: - existing[name] = Queue( + queues.append( + Queue( name, exchange=Exchange(name, type="direct"), routing_key=f"spacedf.tasks.{name}", durable=True, ) + ) routes[f"spacedf.tasks.{name}"] = { "queue": name, "routing_key": f"spacedf.tasks.{name}", } -app.conf.task_queues = tuple(existing.values()) +append_unique_task_queues(app, queues) app.conf.task_routes = routes