UN-3445 [GATED-FEAT] PG-queue auxiliary-worker migration (OSS): loader fix + notification seam

backend/notification_v2/internal_api_views.py

@@ -14,6 +14,7 @@
 from datetime import timedelta
 from typing import Any, cast
 
+from account_v2.models import Organization
 from api_v2.models import APIDeployment
 from django.conf import settings
 from django.db import transaction
@@ -37,6 +38,10 @@
     webhook_url_hash,
 )
 from notification_v2.models import Notification, NotificationBuffer
+from notification_v2.notification_dispatch import (
+    PermanentDispatchError,
+    dispatch_webhook_notification,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -513,6 +518,34 @@ def _reclaim_stale_sending() -> int:
     return int(reclaimed)
 
 
+def _org_identifier(org_pk: int) -> str | None:
+    """Resolve the string ``Organization.organization_id`` from the buffer's org pk.
+
+    ``resolve_transport`` keys its Flipt decision on the org's string identifier,
+    but the buffer stores/uses the Organization pk. One indexed pk lookup per
+    dispatch group (post-commit) — negligible relative to the downstream webhook
+    dispatch.
+
+    Data-anomaly guard: ``NotificationBuffer.organization`` is
+    ``on_delete=CASCADE``, so a live buffer row with a missing org is unreachable
+    in normal operation. A ``None`` here therefore signals a dangling FK — we log
+    it (the only org-traceable breadcrumb; resolve_transport's own warning is keyed
+    on the random dispatch uuid) and fail closed to Celery in resolve_transport.
+    """
+    org_string_id = (
+        Organization.objects.filter(pk=org_pk)
+        .values_list("organization_id", flat=True)
+        .first()
+    )
+    if org_string_id is None:
+        logger.warning(
+            "metric=notification_org_identifier_missing_total org_pk=%s "
+            "(dangling FK; notification routing falls back to Celery)",
+            org_pk,
+        )
+    return org_string_id
+
+
 def _send_clubbed(
     *,
     url: str,
@@ -540,8 +573,12 @@ def _send_clubbed(
     ``buffer_row_ids`` + ``organization_id`` to the worker so it can mark them.
     """
     try:
-        celery_app.send_task(
-            "send_webhook_notification",
+        # Flag-gated transport (UN-3753): PG queue when pg_queue_enabled for this
+        # org, else Celery (byte-identical to the prior send_task). resolve_transport
+        # keys on the org STRING id, but the buffer/worker contract below uses the
+        # org pk — hence _org_identifier(org_id) for routing, org_id in kwargs.
+        dispatch_webhook_notification(
+            celery_app=celery_app,
             args=[url, body, headers, settings.NOTIFICATION_TIMEOUT],
             kwargs={
                 "max_retries": max_retries,
@@ -557,6 +594,7 @@ def _send_clubbed(
                 "organization_id": org_id,
             },
             queue="notifications",
+            org_string_id=_org_identifier(org_id),
         )
         logger.info(
             "metric=notification_batch_dispatched_total platform=%s result=success "
@@ -566,7 +604,34 @@ def _send_clubbed(
             webhook_url_hash(url),
             len(buffer_ids),
         )
+    except PermanentDispatchError:
+        # PG-ONLY permanent failure — enqueue_task validation (priority range /
+        # reply_key+callback exclusivity) or payload serialization — fails
+        # identically every flush tick. Dead-letter now (distinct metric) instead
+        # of reverting to PENDING and re-rendering + re-dispatching + emitting a
+        # broker_failure traceback until the attempt cap. The Celery path never
+        # raises this, so the flag-off flow is unchanged: a Celery send_task failure
+        # is an ordinary Exception handled by the transient branch below, exactly as
+        # before UN-3753. Guard on SENDING so a row the worker already resolved
+        # isn't clobbered.
+        logger.exception(
+            "metric=notification_batch_dispatched_total platform=%s "
+            "result=dispatch_error org_id=%s webhook_url_hash=%s rows=%d",
+            platform,
+            org_id,
+            webhook_url_hash(url),
+            len(buffer_ids),
+        )
+        NotificationBuffer.objects.filter(
+            id__in=buffer_ids,
+            status=BufferStatus.SENDING.value,
+        ).update(status=BufferStatus.DEAD_LETTER.value)
     except Exception:
+        # TRANSIENT transport/broker failure — revert to PENDING (outside the
+        # committed txn) so the next flush tick retries; refund the SENDING-claim
+        # attempt since nothing was queued or sent. Guard on SENDING so a row the
+        # worker already marked terminal (broker raised post-delivery) isn't
+        # resurrected into a duplicate.
         logger.exception(
             "metric=notification_batch_dispatched_total platform=%s "
             "result=broker_failure org_id=%s webhook_url_hash=%s rows=%d",
@@ -575,10 +640,6 @@ def _send_clubbed(
             webhook_url_hash(url),
             len(buffer_ids),
         )
-        # Revert to PENDING (outside the committed txn) so a transient broker
-        # outage retries next tick; refund the SENDING-claim attempt since nothing
-        # was queued or sent. Guard on SENDING so a row the worker already marked
-        # terminal (broker raised post-delivery) isn't resurrected into a duplicate.
         NotificationBuffer.objects.filter(
             id__in=buffer_ids,
             status=BufferStatus.SENDING.value,

backend/notification_v2/notification_dispatch.py

@@ -0,0 +1,119 @@
+"""Transport-routed dispatch for buffered webhook notifications (UN-3753).
+
+Routes the ``send_webhook_notification`` task through the same
+:func:`resolve_transport` flag as the execution path: the PG queue when
+``pg_queue_enabled`` for this org, else Celery. **Fail-closed** — with the gate
+off (the production default) it resolves to Celery, behaving exactly like the
+prior unconditional ``celery_app.send_task`` (zero regression). On PG the task
+is drained by the PG notification consumer on the ``notifications`` queue.
+
+The same ``args``/``kwargs``/``queue`` are forwarded on both paths; on PG they're
+JSON-normalized by ``enqueue_task`` (UUIDs/datetimes → str), so the consumer sees
+the same payload the Celery worker would.
+"""
+
+from __future__ import annotations
+
+import logging
+import uuid
+from typing import Any
+
+from pg_queue.producer import enqueue_task
+from workflow_manager.workflow_v2.transport import resolve_transport
+
+from unstract.core.data_models import is_pg_transport
+
+logger = logging.getLogger(__name__)
+
+# The fired task name — mirrors the Celery task registered by the notification
+# worker; kept as a local constant so the backend doesn't import the workers pkg.
+WEBHOOK_NOTIFICATION_TASK = "send_webhook_notification"
+
+
+class PermanentDispatchError(Exception):
+    """A dispatch failure that would fail identically on every retry.
+
+    Raised ONLY on the PG path, when ``enqueue_task`` rejects the message for a
+    permanent reason (priority range / reply_key+callback exclusivity validation,
+    or a payload that can't be JSON-serialized). The Celery path never raises it,
+    so a caller can dead-letter on this exception without altering the flag-off
+    (Celery) error flow — a Celery ``send_task`` failure stays an ordinary
+    ``Exception`` the caller's transient handler owns, exactly as before.
+    """
+
+
+def dispatch_webhook_notification(
+    *,
+    celery_app: Any,
+    args: list[Any],
+    kwargs: dict[str, Any],
+    queue: str,
+    org_string_id: str | None,
+) -> str:
+    """Dispatch ``send_webhook_notification`` on the resolved transport.
+
+    ``args``/``kwargs``/``queue`` are forwarded unchanged on both paths, so the
+    flag-off (Celery) path is byte-identical to the legacy ``send_task`` call.
+
+    Args:
+        celery_app: Injected Celery app (the backend's ``celery_service.app``);
+            passed in rather than imported so this seam stays trivially testable.
+        args: Positional task args, forwarded verbatim.
+        kwargs: Keyword task args, forwarded verbatim. For the buffered path this
+            carries ``organization_id`` = the buffer's org **pk** (the worker's
+            buffer-mark contract) — deliberately a DIFFERENT identifier from the
+            ``org_string_id`` param below (the two must not be conflated).
+        queue: Target queue name, forwarded verbatim.
+        org_string_id: The org's **string** identifier
+            (``Organization.organization_id``), used solely for the Flipt
+            transport decision — NOT the org pk carried in ``kwargs``. ``None`` (or
+            empty) fails closed to Celery.
+
+    Returns:
+        A task id string — the Celery ``AsyncResult`` id on the Celery path, or
+        the minted PG task id on the PG path. Returned for symmetry / any future
+        caller; the sole current caller (fire-and-forget) discards it.
+    """
+    # A buffered notification is a single fire-and-forget task with no natural
+    # sticky entity, so mint a fresh id to drive Flipt's percentage bucketing and
+    # to serve as the PG task id.
+    dispatch_id = str(uuid.uuid4())
+    # resolve_transport already normalizes falsy input to Celery, so pass the id
+    # straight through (no `or None` needed).
+    transport = resolve_transport(
+        execution_id=dispatch_id,
+        organization_id=org_string_id,
+    )
+    # Use the shared is_pg_transport() — the single source for "what counts as PG
+    # transport" — rather than opening a second comparison site.
+    if is_pg_transport(transport):
+        try:
+            enqueue_task(
+                task_name=WEBHOOK_NOTIFICATION_TASK,
+                queue=queue,
+                args=args,
+                kwargs=kwargs,
+                # enqueue_task's org_id is str-typed — coerce None→"" (unlike the
+                # routing arg above, this `or ""` is load-bearing).
+                org_id=org_string_id or "",
+                task_id=dispatch_id,
+            )
+        except (ValueError, TypeError) as exc:
+            # PG-only permanent failure (enqueue_task validation / JSON encode):
+            # re-raise as PermanentDispatchError so the caller dead-letters it.
+            # A transient PG error (DB down) is NOT wrapped — it propagates as an
+            # ordinary Exception into the caller's retry (revert-to-PENDING) path.
+            raise PermanentDispatchError(str(exc)) from exc
+        logger.info(
+            "Webhook notification enqueued on PG '%s' queue (task_id=%s)",
+            queue,
+            dispatch_id,
+        )
+        return dispatch_id
+    result = celery_app.send_task(
+        WEBHOOK_NOTIFICATION_TASK,
+        args=args,
+        kwargs=kwargs,
+        queue=queue,
+    )
+    return result.id

backend/notification_v2/tests/test_notification_dispatch.py

@@ -0,0 +1,136 @@
+"""Unit tests for the buffered-webhook transport routing (UN-3753).
+
+``resolve_transport`` + ``enqueue_task`` are patched on the module, so no Flipt /
+DB is needed — these pin the routing contract: PG when the flag resolves PG,
+Celery otherwise (fail-closed), with byte-identical args/kwargs/queue on both
+paths.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+import notification_v2.notification_dispatch as nd
+
+_ARGS = ["https://hook.test", {"text": "hi"}, {"Content-Type": "application/json"}, 30]
+_KWARGS = {
+    "max_retries": 3,
+    "retry_delay": 10,
+    "platform": "SLACK",
+    "raise_on_final_failure": True,
+    "buffer_row_ids": ["b1", "b2"],
+    "organization_id": 7,  # the org pk (worker's buffer-mark contract)
+}
+_QUEUE = "notifications"
+
+
+def _dispatch(celery_app, org_string_id="org_x"):
+    return nd.dispatch_webhook_notification(
+        celery_app=celery_app,
+        args=_ARGS,
+        kwargs=_KWARGS,
+        queue=_QUEUE,
+        org_string_id=org_string_id,
+    )
+
+
+class TestDispatchWebhookNotification:
+    def test_routes_to_pg_when_flag_resolves_pg(self):
+        celery = MagicMock()
+        with (
+            patch.object(nd, "resolve_transport", return_value="pg_queue"),
+            patch.object(nd, "enqueue_task", return_value=42) as enqueue,
+        ):
+            task_id = _dispatch(celery)
+        enqueue.assert_called_once()
+        kwargs = enqueue.call_args.kwargs
+        assert kwargs["task_name"] == "send_webhook_notification"
+        assert kwargs["queue"] == _QUEUE
+        assert kwargs["args"] == _ARGS
+        assert kwargs["kwargs"] == _KWARGS
+        assert kwargs["org_id"] == "org_x"
+        # The minted PG task id is returned and threaded into the enqueue row.
+        assert kwargs["task_id"] == task_id
+        celery.send_task.assert_not_called()
+
+    def test_routes_to_celery_when_flag_resolves_celery(self):
+        celery = MagicMock()
+        with (
+            patch.object(nd, "resolve_transport", return_value="celery"),
+            patch.object(nd, "enqueue_task") as enqueue,
+        ):
+            result = _dispatch(celery)
+        celery.send_task.assert_called_once_with(
+            "send_webhook_notification", args=_ARGS, kwargs=_KWARGS, queue=_QUEUE
+        )
+        enqueue.assert_not_called()
+        assert result is celery.send_task.return_value.id
+
+    def test_pg_enqueue_failure_propagates_with_no_celery_fallback(self):
+        # No silent Celery fallback on a TRANSIENT PG failure — it propagates raw
+        # (not wrapped), so the caller (_send_clubbed) reverts rows to PENDING.
+        celery = MagicMock()
+        with (
+            patch.object(nd, "resolve_transport", return_value="pg_queue"),
+            patch.object(nd, "enqueue_task", side_effect=RuntimeError("pg down")),
+        ):
+            with pytest.raises(RuntimeError, match="pg down"):
+                _dispatch(celery)
+        celery.send_task.assert_not_called()
+
+    def test_pg_permanent_enqueue_error_is_wrapped(self):
+        # A PERMANENT enqueue error (ValueError/TypeError: validation / JSON encode)
+        # is re-raised as PermanentDispatchError so the caller dead-letters it —
+        # raised ONLY on the PG path, so the Celery flow is never affected.
+        celery = MagicMock()
+        for exc in (ValueError("priority out of range"), TypeError("not serializable")):
+            with (
+                patch.object(nd, "resolve_transport", return_value="pg_queue"),
+                patch.object(nd, "enqueue_task", side_effect=exc),
+            ):
+                with pytest.raises(nd.PermanentDispatchError):
+                    _dispatch(celery)
+            celery.send_task.assert_not_called()
+
+    def test_none_org_fails_closed_to_celery(self):
+        # A missing org string (org deleted) must not route to PG — resolve_transport
+        # fails closed, and we pass organization_id=None straight through to it.
+        celery = MagicMock()
+        with (
+            patch.object(nd, "resolve_transport", return_value="celery") as resolve,
+            patch.object(nd, "enqueue_task") as enqueue,
+        ):
+            _dispatch(celery, org_string_id=None)
+        assert resolve.call_args.kwargs["organization_id"] is None
+        celery.send_task.assert_called_once()
+        enqueue.assert_not_called()
+
+    def test_resolve_transport_buckets_by_minted_dispatch_id(self):
+        # Fire-and-forget: entity_id is a freshly minted uuid (str), and it equals
+        # the PG task_id so the row and the Flipt bucket agree.
+        celery = MagicMock()
+        with (
+            patch.object(nd, "resolve_transport", return_value="pg_queue") as resolve,
+            patch.object(nd, "enqueue_task", return_value=1) as enqueue,
+        ):
+            task_id = _dispatch(celery)
+        assert resolve.call_args.kwargs["execution_id"] == task_id
+        assert enqueue.call_args.kwargs["task_id"] == task_id
+
+    def test_args_and_kwargs_identical_on_both_paths(self):
+        # The consumer must behave the same regardless of transport.
+        celery = MagicMock()
+        with patch.object(nd, "resolve_transport", return_value="celery"):
+            _dispatch(celery)
+        celery_call = celery.send_task.call_args.kwargs
+        celery2 = MagicMock()
+        with (
+            patch.object(nd, "resolve_transport", return_value="pg_queue"),
+            patch.object(nd, "enqueue_task", return_value=1) as enqueue,
+        ):
+            _dispatch(celery2)
+        assert enqueue.call_args.kwargs["args"] == celery_call["args"]
+        assert enqueue.call_args.kwargs["kwargs"] == celery_call["kwargs"]
+        assert enqueue.call_args.kwargs["queue"] == celery_call["queue"]