pghistory: pass on context to celery tasks

Author: valentijnscholten

dojo/celery.py

@@ -2,7 +2,7 @@
 import os
 from logging.config import dictConfig
 
-from celery import Celery
+from celery import Celery, Task
 from celery.signals import setup_logging
 from django.conf import settings
 
@@ -11,7 +11,31 @@
 # set the default Django settings module for the 'celery' program.
 os.environ.setdefault("DJANGO_SETTINGS_MODULE", "dojo.settings.settings")
 
-app = Celery("dojo")
+
+class PgHistoryTask(Task):
+
+    """
+    Custom Celery base task that automatically applies pghistory context.
+
+    When a task is dispatched via dojo_async_task, the current pghistory
+    context is captured and passed in kwargs as "_pgh_context". This base
+    class extracts that context and applies it before running the task,
+    ensuring all database events share the same context as the original
+    request.
+    """
+
+    def __call__(self, *args, **kwargs):
+        # Import here to avoid circular imports during Celery startup
+        from dojo.pghistory_utils import get_pghistory_context_manager  # noqa: PLC0415
+
+        # Extract context from kwargs (won't be passed to task function)
+        pgh_context = kwargs.pop("_pgh_context", None)
+
+        with get_pghistory_context_manager(pgh_context):
+            return super().__call__(*args, **kwargs)
+
+
+app = Celery("dojo", task_cls=PgHistoryTask)
 
 # Using a string here means the worker will not have to
 # pickle the object when using Windows.

dojo/decorators.py

@@ -83,10 +83,17 @@ def dojo_async_task(func=None, *, signature=False):
     def decorator(func):
         @wraps(func)
         def __wrapper__(*args, **kwargs):
+            from dojo.pghistory_utils import get_serializable_pghistory_context  # noqa: PLC0415 circular import
             from dojo.utils import get_current_user  # noqa: PLC0415 circular import
+
             user = get_current_user()
             kwargs["async_user"] = user
 
+            # Capture pghistory context to pass to Celery worker
+            # The PgHistoryTask base class will apply this context in the worker
+            if pgh_context := get_serializable_pghistory_context():
+                kwargs["_pgh_context"] = pgh_context
+
             dojo_async_task_counter.incr(
                 func.__name__,
                 args=args,

dojo/pghistory_utils.py

@@ -0,0 +1,99 @@
+"""
+Utilities for passing pghistory context to Celery tasks.
+
+pghistory uses thread-local storage, so context is lost when tasks run
+in Celery workers. These utilities allow capturing context in the sender
+process and recreating it in the worker.
+"""
+import uuid
+from contextlib import nullcontext
+
+from pghistory import runtime as pghistory_runtime
+
+
+def get_serializable_pghistory_context():
+    """
+    Capture the current pghistory context for passing to Celery tasks.
+
+    Returns a JSON-serializable dict with context id and metadata,
+    or None if no context is active.
+    """
+    if hasattr(pghistory_runtime._tracker, "value"):
+        ctx = pghistory_runtime._tracker.value
+        return {
+            "id": str(ctx.id),
+            "metadata": ctx.metadata.copy(),
+        }
+    return None
+
+
+class PgHistoryContextFromTask:
+
+    """
+    Context manager to apply pghistory context received from a Celery task.
+
+    This recreates the exact same context (with the same UUID) that was
+    active when the task was dispatched, ensuring all events share the
+    same pgh_context_id.
+
+    Usage:
+        pgh_context = kwargs.pop("_pgh_context", None)
+        with PgHistoryContextFromTask(pgh_context):
+            # Task body runs here with context applied
+    """
+
+    def __init__(self, context_data):
+        """
+        Initialize with context data from Celery kwargs.
+
+        Args:
+            context_data: Dict with "id" (UUID string) and "metadata" (dict),
+                         or None for no-op behavior.
+
+        """
+        self.context_data = context_data
+        self._pre_execute_hook = None
+        self._owns_context = False
+
+    def __enter__(self):
+        if not self.context_data:
+            return None
+
+        from django.db import connection  # noqa: PLC0415
+
+        context_id = uuid.UUID(self.context_data["id"])
+        metadata = self.context_data["metadata"]
+
+        # Only create a new context if one doesn't already exist
+        if not hasattr(pghistory_runtime._tracker, "value"):
+            self._pre_execute_hook = connection.execute_wrapper(
+                pghistory_runtime._inject_history_context,
+            )
+            self._pre_execute_hook.__enter__()
+            pghistory_runtime._tracker.value = pghistory_runtime.Context(
+                id=context_id,
+                metadata=metadata,
+            )
+            self._owns_context = True
+        else:
+            # Context already exists, just merge metadata
+            pghistory_runtime._tracker.value.metadata.update(metadata)
+
+        return pghistory_runtime._tracker.value
+
+    def __exit__(self, *exc):
+        if self._owns_context and self._pre_execute_hook:
+            delattr(pghistory_runtime._tracker, "value")
+            self._pre_execute_hook.__exit__(*exc)
+
+
+def get_pghistory_context_manager(context_data):
+    """
+    Return appropriate context manager for the given context data.
+
+    Returns PgHistoryContextFromTask if context_data is provided,
+    otherwise returns a no-op nullcontext.
+    """
+    if context_data:
+        return PgHistoryContextFromTask(context_data)
+    return nullcontext()