From 33e9268ce5498b926744bc9c6e9674c1569d04ca Mon Sep 17 00:00:00 2001 From: lethanhdat762003 Date: Wed, 5 Aug 2026 14:58:23 +0700 Subject: [PATCH] feat: add shared subscription routing utilities --- common/apps/billing/constants.py | 1 + common/celery/constants.py | 10 ++++++ common/celery/routing.py | 59 +++++++++++++++++++++++++++----- common/celery/task_senders.py | 22 ++++++++++++ common/celery/tasks.py | 8 +++++ 5 files changed, 91 insertions(+), 9 deletions(-) diff --git a/common/apps/billing/constants.py b/common/apps/billing/constants.py index 3008dd9..de24ff0 100644 --- a/common/apps/billing/constants.py +++ b/common/apps/billing/constants.py @@ -10,3 +10,4 @@ class FeatureCode: DASHBOARD_MAX_COUNT = "dashboard.max_count" DASHBOARD_UPDATE_COUNT = "dashboard.update_count" WHITELABEL_ENABLED = "whitelabel.enabled" + AUTOMATION_MAX_COUNT = "automation.max_count" diff --git a/common/celery/constants.py b/common/celery/constants.py index 672fdc4..48f28e0 100644 --- a/common/celery/constants.py +++ b/common/celery/constants.py @@ -6,3 +6,13 @@ CONSOLE_SERVICE = "console_service" CONSOLE_SERVICE_ADD_OR_REMOVE_SPACE = f"{CONSOLE_SERVICE}.add_or_remove_space" CONSOLE_SERVICE_DELETE_UPLOAD_FILE = f"{CONSOLE_SERVICE}.delete_upload_file" + +SUBSCRIPTION_DOWNGRADE_EXCHANGE = "subscription_downgrade" +SUBSCRIPTION_UPGRADE_EXCHANGE = "subscription_upgrade" +SUBSCRIPTION_LIFECYCLE_EXCHANGES = { + "downgrade": SUBSCRIPTION_DOWNGRADE_EXCHANGE, + "upgrade": SUBSCRIPTION_UPGRADE_EXCHANGE, +} + +TELEMETRY_SUBSCRIPTION_EVENT_AUTOMATION = "automation" +TELEMETRY_SUBSCRIPTION_EVENT_ENTITIES = "entities" diff --git a/common/celery/routing.py b/common/celery/routing.py index fe8d75e..c8cbf31 100644 --- a/common/celery/routing.py +++ b/common/celery/routing.py @@ -2,6 +2,33 @@ from django.utils.module_loading import import_string from kombu import Exchange, Queue +from common.celery.constants import SUBSCRIPTION_LIFECYCLE_EXCHANGES + + +def get_queue_name(queue): + name = getattr(queue, "name", queue) + if isinstance(name, dict): + name = name.get("name") or name.get("queue") + return name + + +def append_unique_task_queues(celery_app, queues): + if celery_app.conf.task_queues is None: + celery_app.conf.task_queues = () + + existing = { + name: queue + for queue in celery_app.conf.task_queues + if (name := get_queue_name(queue)) + } + + for queue in queues: + name = get_queue_name(queue) + if name and name not in existing: + existing[name] = queue + + celery_app.conf.task_queues = tuple(existing.values()) + def setup_synchronous_model_task_routing(): celery_app = import_string(settings.CELERY_APP) @@ -79,12 +106,16 @@ def setup_organization_task_routing(): "queue": queue_name, "routing_key": queue_cfg["routing_key"], } - celery_app.conf.task_queues = celery_app.conf.task_queues + tuple(new_queues) + append_unique_task_queues(celery_app, new_queues) -def setup_subscription_task_routing(task_names): - """ - Setup task routing for subscription downgrade and upgrade tasks. +def setup_subscription_task_routing(task_specs): + """Setup routing for subscription lifecycle tasks. + + Each task spec must provide: + - task_name: celery task name without the ``spacedf.tasks.`` prefix + - service: downstream service routing prefix, e.g. ``telemetry`` + - lifecycle: ``downgrade`` or ``upgrade`` """ celery_app = import_string(settings.CELERY_APP) @@ -94,17 +125,27 @@ def setup_subscription_task_routing(task_names): celery_app.conf.task_routes = {} new_queues = [] - for name in task_names: + for task_spec in task_specs: + name = task_spec["task_name"] + service_name = task_spec["service"] + lifecycle = task_spec["lifecycle"] + + exchange_name = SUBSCRIPTION_LIFECYCLE_EXCHANGES.get(lifecycle) + if exchange_name is None: + raise ValueError(f"unsupported subscription lifecycle: {lifecycle}") + + routing_key = f"{service_name}.{lifecycle}" new_queues.append( Queue( name, - exchange=Exchange(name, type="direct"), - routing_key=f"spacedf.tasks.{name}", + exchange=Exchange(exchange_name, type="direct"), + routing_key=routing_key, + durable=True, ) ) celery_app.conf.task_routes[f"spacedf.tasks.{name}"] = { "queue": name, - "routing_key": f"spacedf.tasks.{name}", + "routing_key": routing_key, } - celery_app.conf.task_queues = celery_app.conf.task_queues + tuple(new_queues) + append_unique_task_queues(celery_app, new_queues) diff --git a/common/celery/task_senders.py b/common/celery/task_senders.py index ec2538e..66c1f1d 100644 --- a/common/celery/task_senders.py +++ b/common/celery/task_senders.py @@ -1,6 +1,8 @@ from django.conf import settings from django.utils.module_loading import import_string +from common.celery.constants import SUBSCRIPTION_LIFECYCLE_EXCHANGES + def send_task(name, message, **kwargs): celery_app = import_string(settings.CELERY_APP) @@ -15,3 +17,23 @@ def send_task(name, message, **kwargs): kwargs=message, **kwargs, ) + + +def send_subscription_task(service, lifecycle, task_name, message, **kwargs): + celery_app = import_string(settings.CELERY_APP) + exchange = SUBSCRIPTION_LIFECYCLE_EXCHANGES.get(lifecycle) + if exchange is None: + raise ValueError(f"unsupported subscription lifecycle: {lifecycle}") + + routing_key = f"{service}.{lifecycle}" + return celery_app.send_task( + name=f"spacedf.tasks.{task_name}", + exchange=exchange, + routing_key=routing_key, + retry=True, + retry_policy=dict( + max_retries=3, interval_start=3, interval_step=1, interval_max=6 + ), + kwargs=message, + **kwargs, + ) diff --git a/common/celery/tasks.py b/common/celery/tasks.py index 1dc1807..157484c 100644 --- a/common/celery/tasks.py +++ b/common/celery/tasks.py @@ -1,12 +1,17 @@ import logging from celery import shared_task +from celery.exceptions import Ignore from django.db import models, transaction from django_tenants.utils import get_tenant_model, tenant_context logger = logging.getLogger(__name__) +class PermanentTaskError(Exception): + """Signals an invalid task payload that should not be retried.""" + + def task(max_retries=3, bind=False, task_acks_late=True, prefetch_count=1, **option): """ Use as a decorator for the Celery shared_task @@ -28,6 +33,9 @@ def inner(self, **kwargs): handler(self=self, **kwargs) else: handler(**kwargs) + except PermanentTaskError as e: + logger.exception(e) + raise Ignore() except Exception as e: logger.exception(e) self.retry(countdown=3**self.request.retries)