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,12 +38,25 @@
     webhook_url_hash,
 )
 from notification_v2.models import Notification, NotificationBuffer
+from notification_v2.notification_dispatch import (
+    PermanentDispatchError,
+    dispatch_webhook_notification,
+)
 
 logger = logging.getLogger(__name__)
 
 # Constants for error messages
 INTERNAL_SERVER_ERROR_MSG = "Internal server error"
 
+# How many early transient-failure cycles get their SENDING claim refunded before the
+# claim starts sticking. Refunding forever makes net attempt progress zero, so
+# ``NOTIFICATION_MAX_DISPATCH_ATTEMPTS`` becomes unreachable and a permanently
+# recurring failure (not just a brief broker blip) re-dispatches every flush tick with
+# no termination condition. Small by design: a genuine blip clears well inside this,
+# and anything still failing after it is treated as repetitive and allowed to age
+# toward the cap.
+NOTIFICATION_TRANSIENT_REFUND_LIMIT = 3
+
 
 def _load_execution(execution_id: str | None) -> WorkflowExecution | None:
     """Best-effort lookup; returns None on missing id or unknown row."""
@@ -513,6 +527,38 @@ 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:
+        # Sentry-routed (logger.error): a live buffer row with no org is a data
+        # anomaly (dangling FK / corruption) that shouldn't happen under the
+        # CASCADE constraint, not routine noise. Routing still fails closed to
+        # Celery in resolve_transport.
+        logger.error(
+            "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,
@@ -521,7 +567,7 @@ def _send_clubbed(
     platform: str,
     max_retries: int,
     buffer_ids: list[str],
-    org_id: Any,
+    org_id: int,
 ) -> None:
     """Send the clubbed Celery task after the DB transition has committed.
 
@@ -540,8 +586,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.
+        dispatched = dispatch_webhook_notification(
+            celery_app=celery_app,
             args=[url, body, headers, settings.NOTIFICATION_TIMEOUT],
             kwargs={
                 "max_retries": max_retries,
@@ -557,16 +607,61 @@ def _send_clubbed(
                 "organization_id": org_id,
             },
             queue="notifications",
+            org_string_id=_org_identifier(org_id),
         )
+        # transport= makes the rollout answerable from the logs: during a percentage
+        # ramp the question is "are PG-routed notifications succeeding at the same
+        # rate as Celery-routed ones?", which result=success alone cannot answer.
         logger.info(
             "metric=notification_batch_dispatched_total platform=%s result=success "
-            "org_id=%s webhook_url_hash=%s rows=%d",
+            "transport=%s org_id=%s webhook_url_hash=%s rows=%d task_id=%s",
             platform,
+            dispatched.transport,
             org_id,
             webhook_url_hash(url),
             len(buffer_ids),
+            dispatched.task_id,
         )
+    except PermanentDispatchError:
+        # PG-ONLY permanent failure. From this path the only reachable cause is
+        # payload JSON-serialization (e.g. a NaN/Infinity float that jsonb rejects
+        # at insert): this call passes no reply_key/callback and a default in-range
+        # priority, so enqueue_task's other permanent checks can't fire here (the
+        # full set lives on PermanentDispatchError's docstring). It 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. Guard on SENDING so a row
+        # the worker already marked terminal (broker raised post-delivery) isn't
+        # resurrected into a duplicate.
+        #
+        # The refund is BOUNDED, not unconditional: refunding the SENDING claim on
+        # every cycle makes net progress zero, so `NOTIFICATION_MAX_DISPATCH_ATTEMPTS`
+        # can never be reached and a permanently-recurring failure re-renders and
+        # re-dispatches forever (emitting a traceback each tick). Refunding is still
+        # correct for a genuinely transient blip — nothing was queued or sent — so we
+        # keep it while attempts are low and let the cap take over once a failure has
+        # proven itself repetitive. This matters more now that this branch also
+        # catches non-permanent PG enqueue failures (e.g. a DataError that will recur
+        # identically), not just "RabbitMQ is briefly down".
         logger.exception(
             "metric=notification_batch_dispatched_total platform=%s "
             "result=broker_failure org_id=%s webhook_url_hash=%s rows=%d",
@@ -575,21 +670,24 @@ 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(
+        reverted = NotificationBuffer.objects.filter(
             id__in=buffer_ids,
             status=BufferStatus.SENDING.value,
-        ).update(
+        )
+        # Refund only the early attempts; past the limit the claim stands so the
+        # counter grows and the cap can eventually dead-letter the group.
+        reverted.filter(
+            dispatch_attempts__lte=NOTIFICATION_TRANSIENT_REFUND_LIMIT
+        ).update(dispatch_attempts=F("dispatch_attempts") - 1)
+        # Every claimed row goes back to PENDING for the next flush tick, refunded
+        # or not.
+        reverted.update(
             status=BufferStatus.PENDING.value,
             dispatched_at=None,
-            dispatch_attempts=F("dispatch_attempts") - 1,
         )
 
 
-def _penalize_render_failure(buffer_ids: list[str], org_id: Any, platform: str) -> None:
+def _penalize_render_failure(buffer_ids: list[str], org_id: int, platform: str) -> None:
     """Charge a dispatch attempt to a group whose payloads failed to render.
 
     The SENDING-claim increment never runs on a render failure, so count it here
@@ -609,7 +707,7 @@ def _penalize_render_failure(buffer_ids: list[str], org_id: Any, platform: str)
 
 
 def _dispatch_group(
-    org_id: Any,
+    org_id: int,
     webhook_url: str,
     auth_sig: str,
     platform: str,

backend/notification_v2/notification_dispatch.py

@@ -0,0 +1,180 @@
+"""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).
+
+DEPLOYMENT PREREQUISITE (PG path): on PG the task lands on the ``notifications``
+queue and requires a pg-queue consumer configured with
+``WORKER_PG_QUEUE_CONSUMER_WORKER_TYPE=notification`` /
+``WORKER_PG_QUEUE_CONSUMER_QUEUE=notifications``. **No such service exists in this
+repo's compose yet** — none of the ``pg-queue-consumer`` services in
+``docker/docker-compose.yaml`` polls ``notifications``, and ``run-worker.sh``'s
+``PG_CONSUMER_ROLES`` has no notification role. Until one is deployed the flag must
+stay off for the org: enqueued rows would sit undrained (no TTL sweep covers
+``pg_queue_message``) and every buffered webhook for that org would be lost.
+
+``args`` and ``queue`` are forwarded verbatim on both paths. ``kwargs`` differ in
+exactly two keys, forced on the PG branch only (see the inline note at the branch):
+``raise_on_final_failure`` → ``False`` and ``max_retries`` → ``0``. On PG the
+payload is additionally JSON-normalized by ``enqueue_task`` (UUIDs/datetimes → str).
+"""
+
+from __future__ import annotations
+
+import logging
+import uuid
+from typing import Any, NamedTuple
+
+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"
+
+
+# Transport labels for the dispatch metric. During a percentage ramp the one thing
+# you need from the logs is "are PG-routed notifications succeeding at the same rate
+# as Celery-routed ones?" — a bare task id can't answer that, since a minted PG id
+# and a Celery AsyncResult id are indistinguishable to the caller.
+PG_TRANSPORT = "pg_queue"
+CELERY_TRANSPORT = "celery"
+
+
+class DispatchResult(NamedTuple):
+    """Which transport actually took the dispatch, plus the resulting task id."""
+
+    transport: str
+    task_id: str
+
+
+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,
+) -> DispatchResult:
+    """Dispatch ``send_webhook_notification`` on the resolved transport.
+
+    ``args``/``queue`` are forwarded unchanged on both paths, and on the Celery
+    branch ``kwargs`` too — so the flag-off path is byte-identical to the legacy
+    ``send_task`` call. The PG branch overrides two retry-semantics kwargs (see
+    below); nothing else differs.
+
+    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 as-is on Celery; on PG,
+            ``raise_on_final_failure`` is overridden to ``False`` and
+            ``max_retries`` to ``0`` (the in-task retry loop is a no-op under the
+            consumer's eager ``apply()`` — see the inline note). 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 :class:`DispatchResult` carrying which transport took the dispatch and
+        the resulting task id (the Celery ``AsyncResult`` id, or the minted PG task
+        id). The transport is what makes the caller's dispatch metric answerable
+        during a ramp — the two id flavours are otherwise indistinguishable.
+    """
+    # 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):
+        # Two PG-only kwarg overrides, both about the SAME thing: the consumer runs
+        # the task eagerly via ``task.apply(..., throw=True)``, where Celery's
+        # in-task retry loop does not work and its terminal branch is what we need.
+        #
+        # 1. ``max_retries`` → 0. Under ``apply()`` ``self.request.retries`` is ALWAYS
+        #    0, so the worker's ``if self.request.retries < max_retries`` guard
+        #    (workers/notification/tasks.py) is true on the FIRST failure for any
+        #    max_retries >= 1 and calls ``self.retry(...)``, which raises ``Retry``;
+        #    with ``throw=True`` that propagates straight out of ``apply()`` and the
+        #    terminal branch is never reached — so the buffers are never marked
+        #    DEAD_LETTER and ``raise_on_final_failure`` is never even read. The
+        #    consumer's ``except Exception`` then leaves the row for vt-expiry
+        #    redelivery, re-POSTing the subscriber every time. Forcing 0 sends the
+        #    task down the terminal branch on the first failure instead.
+        # 2. ``raise_on_final_failure`` → False. On BOTH transports the worker marks
+        #    the buffers DEAD_LETTER over the internal API *before* it re-raises, so
+        #    the re-raise is only a FAILURE-state signal for Celery monitoring (no
+        #    redelivery there). On the PG consumer that same raise is treated as a
+        #    failure and leaves the row for redelivery, so it must not raise.
+        #
+        # Together: one POST, buffers dead-lettered, task returns None → the consumer
+        # acks (deletes) the row. Retry spacing belongs to the PG layer, not the task.
+        # (The Celery branch below keeps kwargs verbatim — byte-identical.)
+        pg_kwargs = {**kwargs, "raise_on_final_failure": False, "max_retries": 0}
+        try:
+            msg_id = enqueue_task(
+                task_name=WEBHOOK_NOTIFICATION_TASK,
+                queue=queue,
+                args=args,
+                kwargs=pg_kwargs,
+                # enqueue_task types org_id as str; None→"" here satisfies that
+                # type only, not runtime — enqueue_task itself re-coerces
+                # ``org_id or ""`` at insert (producer.py), so this has no runtime
+                # effect.
+                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
+        # msg_id correlates this line with the producer's own ``msg_id=`` log, so a
+        # dropped notification can be traced across the seam/producer boundary.
+        logger.info(
+            "Webhook notification enqueued on PG '%s' queue (task_id=%s msg_id=%s)",
+            queue,
+            dispatch_id,
+            msg_id,
+        )
+        return DispatchResult(transport=PG_TRANSPORT, task_id=dispatch_id)
+    result = celery_app.send_task(
+        WEBHOOK_NOTIFICATION_TASK,
+        args=args,
+        kwargs=kwargs,
+        queue=queue,
+    )
+    return DispatchResult(transport=CELERY_TRANSPORT, task_id=result.id)