initial changes to record args & kwargs

Author: snopoke

taskbadger/celery.py

@@ -18,7 +18,7 @@
 
 KWARG_PREFIX = "taskbadger_"
 TB_KWARGS_ARG = f"{KWARG_PREFIX}kwargs"
-IGNORE_ARGS = {TB_KWARGS_ARG, f"{KWARG_PREFIX}task", f"{KWARG_PREFIX}task_id"}
+IGNORE_ARGS = {TB_KWARGS_ARG, f"{KWARG_PREFIX}task", f"{KWARG_PREFIX}task_id", f"{KWARG_PREFIX}record_task_args"}
 TB_TASK_ID = f"{KWARG_PREFIX}task_id"
 
 TERMINAL_STATES = {
@@ -124,6 +124,8 @@ def apply_async(self, *args, **kwargs):
         if Badger.is_configured():
             headers["taskbadger_track"] = True
             headers[TB_KWARGS_ARG] = tb_kwargs
+            if "record_task_args" in tb_kwargs:
+                headers["taskbadger_record_task_args"] = tb_kwargs.pop("record_task_args")
 
         result = super().apply_async(*args, **kwargs)
 
@@ -187,6 +189,12 @@ def task_publish_handler(sender=None, headers=None, body=None, **kwargs):
     kwargs["status"] = StatusEnum.PENDING
     name = kwargs.pop("name", headers["task"])
 
+    global_record_task_args = celery_system and celery_system.record_task_args
+    if headers.get("taskbadger_record_task_args", global_record_task_args):
+        data = kwargs.setdefault("data", {})
+        data["celery_task_args"] = body[0]
+        data["celery_task_kwargs"] = body[1]
+
     task = create_task_safe(name, **kwargs)
     if task:
         meta = {TB_TASK_ID: task.id}

taskbadger/systems/celery.py

@@ -6,7 +6,7 @@
 class CelerySystemIntegration(System):
     identifier = "celery"
 
-    def __init__(self, auto_track_tasks=True, includes=None, excludes=None):
+    def __init__(self, auto_track_tasks=True, includes=None, excludes=None, record_task_args=False):
         """
         Args:
             auto_track_tasks: Automatically track all Celery tasks regardless of whether they are using the
@@ -16,10 +16,12 @@ def __init__(self, auto_track_tasks=True, includes=None, excludes=None):
                 matches both an include and an exclude, it will be excluded.
             excludes: A list of task names to exclude from tracking. As with `includes`, these can be either
                 the full task name or a regular expression. Exclusions take precedence over inclusions.
+            record_task_args: Record the arguments passed to each task.
         """
         self.auto_track_tasks = auto_track_tasks
         self.includes = includes
         self.excludes = excludes
+        self.record_task_args = record_task_args
 
         if auto_track_tasks:
             # Importing this here ensures that the Celery signal handlers are registered

tests/test_celery.py

@@ -111,6 +111,74 @@ def add_with_task_args(self, a, b):
     create.assert_called_once_with("new_name", value_max=10, actions=actions, status=StatusEnum.PENDING)
 
 
+def test_celery_record_args(celery_session_app, celery_session_worker, bind_settings):
+    @celery_session_app.task(bind=True, base=Task)
+    def add_with_task_args(self, a, b):
+        assert self.taskbadger_task is not None
+        return a + b
+
+    celery_session_worker.reload()
+
+    with (
+        mock.patch("taskbadger.celery.create_task_safe") as create,
+        mock.patch("taskbadger.celery.update_task_safe"),
+        mock.patch("taskbadger.celery.get_task"),
+    ):
+        create.return_value = task_for_test()
+
+        result = add_with_task_args.apply_async(
+            (2, 2),
+            taskbadger_name="new_name",
+            taskbadger_value_max=10,
+            taskbadger_kwargs={"data": {"foo": "bar"}},
+            taskbadger_record_task_args=True,
+        )
+        assert result.get(timeout=10, propagate=True) == 4
+
+    create.assert_called_once_with(
+        "new_name",
+        value_max=10,
+        data={"foo": "bar", "celery_task_args": (2, 2), "celery_task_kwargs": {}},
+        status=StatusEnum.PENDING,
+    )
+
+
+def test_celery_record_task_kwargs(celery_session_app, celery_session_worker, bind_settings):
+    @celery_session_app.task(bind=True, base=Task)
+    def add_with_task_args(self, a, b, c=0):
+        assert self.taskbadger_task is not None
+        return a + b + c
+
+    celery_session_worker.reload()
+
+    with (
+        mock.patch("taskbadger.celery.create_task_safe") as create,
+        mock.patch("taskbadger.celery.update_task_safe"),
+        mock.patch("taskbadger.celery.get_task"),
+    ):
+        create.return_value = task_for_test()
+
+        actions = [Action("stale", integration=EmailIntegration(to="test@test.com"))]
+        result = add_with_task_args.delay(
+            2,
+            2,
+            c=3,
+            taskbadger_name="new_name",
+            taskbadger_value_max=10,
+            taskbadger_kwargs={"actions": actions},
+            taskbadger_record_task_args=True,
+        )
+        assert result.get(timeout=10, propagate=True) == 7
+
+    create.assert_called_once_with(
+        "new_name",
+        value_max=10,
+        data={"celery_task_args": (2, 2), "celery_task_kwargs": {"c": 3}},
+        actions=actions,
+        status=StatusEnum.PENDING,
+    )
+
+
 def test_celery_task_with_args_in_decorator(celery_session_app, celery_session_worker, bind_settings):
     @celery_session_app.task(
         bind=True,