First version of multi-tenant setup (EE)

snappea/decorators.py

@@ -6,6 +6,7 @@ from performance.context_managers import time_to_logger
 from . import registry
 from .models import Task, wakeup_server
 from .settings import get_settings
+from .utils import add_task_kwargs
 
 performance_logger = logging.getLogger("bugsink.performance.snappea")
 
@@ -26,6 +27,7 @@ def shared_task(function):
             # No need for a transaction: we just write something (not connected to any other object, and we will never
             # touch it again). Counterpoint: if we'd have a transaction, we could distinguish between "wait for write
             # lock" and "actually write".
+            kwargs.update(add_task_kwargs())
             Task.objects.create(task_name=name, args=json.dumps(args), kwargs=json.dumps(kwargs))
 
             # not necessary: `connections["snappea"].close()`; Django does this at the end of the request and the

snappea/foreman.py

@@ -22,6 +22,8 @@ from . import registry
 from .models import Task
 from .datastructures import Workers
 from .settings import get_settings
+from .utils import run_task_context
+
 
 logger = logging.getLogger("snappea.foreman")
 performance_logger = logging.getLogger("bugsink.performance.snappea")
@@ -170,7 +172,9 @@ class Foreman:
         def non_failing_function(*inner_args, **inner_kwargs):
             t0 = time.time()
             try:
-                function(*inner_args, **inner_kwargs)
+                with run_task_context(inner_args, inner_kwargs):
+                    function(*inner_args, **inner_kwargs)
+
             except Exception as e:
                 # Potential TODO: make this configurable / depend on our existing config in bugsink/settings.py
                 logger.warning("Snappea caught Exception: %s", str(e))

snappea/settings.py

@@ -32,6 +32,8 @@ DEFAULTS = {
 
     "TASK_QS_LIMIT": 100,
 
+    "HOOK_ADD_TASK_KWARGS": "snappea.utils.dont_add_anything",
+    "HOOK_RUN_TASK_CONTEXT": "snappea.utils.no_context",
 }
 
 

snappea/utils.py

@@ -0,0 +1,41 @@
+from contextlib import contextmanager
+import importlib
+
+from .settings import get_settings
+
+
+def add_task_kwargs():
+    """Hook for extending Task kwargs"""
+
+    if not hasattr(add_task_kwargs, "func"):
+        # the configured function is cached on add_task_kwargs itself
+        hook = get_settings().HOOK_ADD_TASK_KWARGS
+        module_name, function_name = hook.rsplit('.', 1)
+        module = importlib.import_module(module_name)
+        add_task_kwargs.func = getattr(module, function_name)
+
+    return add_task_kwargs.func()
+
+
+def run_task_context(task_args, task_kwargs):
+    """Hook for running a task in a context; the task's args and kwargs are passed for optional pre-processing"""
+
+    if not hasattr(add_task_kwargs, "func"):
+        # the configured function is cached on run_task_context itself
+        hook = get_settings().HOOK_RUN_TASK_CONTEXT
+        module_name, function_name = hook.rsplit('.', 1)
+        module = importlib.import_module(module_name)
+        run_task_context.func = getattr(module, function_name)
+
+    return run_task_context.func(task_args, task_kwargs)
+
+
+def dont_add_anything():
+    # no-op impl of add_task_kwargs
+    return {}
+
+
+@contextmanager
+def no_context(task_args, task_kwargs):
+    # no-op impl of run_task_context
+    yield