feat(procrastinate): record_task_args stores defer kwargs in task data

Author: snopoke

taskbadger/procrastinate.py

@@ -14,6 +14,7 @@
 import collections
 import functools
 import inspect
+import json
 import logging
 from contextvars import ContextVar
 
@@ -215,9 +216,17 @@ def _maybe_create_pending(task, kwargs):
     for key in ("value_max", "tags"):
         if key in opts and opts[key] is not None:
             create_kwargs[key] = opts[key]
-    user_data = opts.get("data")
-    if user_data:
-        create_kwargs["data"] = dict(user_data)
+
+    data = dict(opts.get("data") or {})
+
+    record_args = opts.get("record_task_args")
+    if record_args is None:
+        record_args = bool(system) and system.record_task_args
+    if record_args:
+        data["procrastinate_task_kwargs"] = _serialize_kwargs(kwargs)
+
+    if data:
+        create_kwargs["data"] = data
 
     tb_task = create_task_safe(name, **create_kwargs)
     if tb_task is None:
@@ -228,6 +237,19 @@ def _maybe_create_pending(task, kwargs):
     return new_kwargs
 
 
+def _serialize_kwargs(kwargs):
+    """Return a JSON-roundtrippable copy of the defer kwargs.
+
+    Procrastinate already requires kwargs be JSON-serializable, so a json
+    dumps/loads roundtrip is safe. Non-serializable values are dropped with
+    a warning."""
+    try:
+        return json.loads(json.dumps(kwargs))
+    except (TypeError, ValueError) as e:
+        log.warning("Error serializing task arguments: %s", e)
+        return {}
+
+
 _TRACK_OPT_KEYS = ("name", "value_max", "tags", "data", "record_task_args")
 
 

tests/test_procrastinate.py

@@ -307,3 +307,20 @@ def self_complete():
     assert StatusEnum.PROCESSING in statuses
     # Last attempted SUCCESS call should be suppressed
     assert statuses.count(StatusEnum.SUCCESS) == 0
+
+
+@pytest.mark.usefixtures("_bind_settings")
+def test_record_task_args_stores_kwargs(app):
+    @track(record_task_args=True, data={"existing": 1})
+    @app.task(name="recorder")
+    def recorder(a, b):
+        return a + b
+
+    tb = task_for_test()
+    with mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb) as create:
+        recorder.defer(a=5, b=6)
+
+    assert create.call_args.kwargs["data"] == {
+        "existing": 1,
+        "procrastinate_task_kwargs": {"a": 5, "b": 6},
+    }