[celery#305] Added a sanity check to task_props_extension

diegocastrum · diegocastrum · commit 23b9cd53bb3a · 2022-08-17T08:00:53.000+02:00
Added `get_task_props_extension` to `settings` module which will raise an `ImproperlyConfigured` when the task_props_extension doesn't complies with the Mapping protocol. `DatabaseBackend` will make use of `get_task_props_extension` to update a potential custom model with the custom properties --- Resolves celery#305 Fixes celery#314
diff --git a/django_celery_results/backends/database.py b/django_celery_results/backends/database.py
@@ -1,5 +1,6 @@
 import binascii
 import json
+from typing import Mapping
 
 from celery import maybe_signature
 from celery.backends.base import BaseDictBackend
@@ -13,7 +14,7 @@
 
 from ..models import ChordCounter
 from ..models.helpers import taskresult_model, groupresult_model
-from ..settings import extend_task_props_callback
+from ..settings import get_task_props_extension
 
 EXCEPTIONS_TO_CATCH = (InterfaceError,)
 
@@ -125,14 +126,8 @@ def _store_result(
             'using': using,
         }
 
-        task_props.update(
-            self._get_extended_properties(request, traceback)
-        )
-
-        # TODO: Wrap this and make some sanity checks to complain the Mapping
-        # protocol.
-        task_props.update(
-            extend_task_props_callback(request, dict(task_props)))
+        task_props.update(self._get_extended_properties(request, traceback))
+        task_props.update(get_task_props_extension(request, dict(task_props)))
 
         self.TaskModel._default_manager.store_result(**task_props)
         return result
diff --git a/django_celery_results/settings.py b/django_celery_results/settings.py
@@ -1,5 +1,6 @@
 from django.conf import settings
 from django.core.exceptions import ImproperlyConfigured
+from collections.abc import Mapping
 
 
 def get_callback_function(settings_name, default=None):
@@ -16,5 +17,21 @@ def get_callback_function(settings_name, default=None):
 
 
 extend_task_props_callback = get_callback_function(
-    "CELERY_RESULTS_EXTEND_TASK_PROPS_CALLBACK", dict
+    "CELERY_RESULTS_EXTEND_TASK_PROPS_CALLBACK"
 )
+
+
+def get_task_props_extension(request, task_props):
+    """Extend the task properties with custom properties to fill custom models."""
+
+    task_props_extension = extend_task_props_callback(request, task_props) or {}
+    if task_props_extension is None:
+        return {}
+
+    if not isinstance(task_props_extension, Mapping):
+        raise ImproperlyConfigured(
+            "CELERY_RESULTS_EXTEND_TASK_PROPS_CALLBACK must return a Mapping "
+            "instance."
+        )
+
+    return task_props_extension
