support custom serialization

Author: snopoke

taskbadger/celery.py

@@ -11,6 +11,7 @@
     task_retry,
     task_success,
 )
+from kombu import serialization
 
 from .internal.models import StatusEnum
 from .mug import Badger
@@ -197,7 +198,8 @@ def task_publish_handler(sender=None, headers=None, body=None, **kwargs):
             "celery_task_kwargs": body[1],
         }
         try:
-            json.dumps(data)
+            _, _, value = serialization.dumps(data, serializer="json")
+            data = json.loads(value)
         except Exception:
             log.error("Error serializing task arguments for task '%s'", name)
         else:

tests/test_celery.py

@@ -13,6 +13,7 @@
 
 import celery
 import pytest
+from kombu.utils.json import register_type
 
 from taskbadger import Action, EmailIntegration, StatusEnum
 from taskbadger.celery import Task
@@ -179,6 +180,41 @@ def add_with_task_args(self, a, b, c=0):
     )
 
 
+def test_celery_record_task_args_custom_serialization(celery_session_app, celery_session_worker, bind_settings):
+    class A:
+        def __init__(self, a, b):
+            self.a = a
+            self.b = b
+
+    register_type(A, "A", lambda o: [o.a, o.b], lambda o: A(*o))
+
+    @celery_session_app.task(bind=True, base=Task)
+    def add_with_task_args(self, a):
+        assert self.taskbadger_task is not None
+        return a.a + 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.delay(
+            A(2, 2),
+            taskbadger_record_task_args=True,
+        )
+        assert result.get(timeout=10, propagate=True) == 4
+
+    create.assert_called_once_with(
+        "tests.test_celery.add_with_task_args",
+        data={"celery_task_args": [{"__type__": "A", "__value__": [2, 2]}], "celery_task_kwargs": {}},
+        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,