feat: add shared subscription routing utilities

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"

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"

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)

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,
+    )

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)