[celery] Drop celery.RequestIDAwareTask in favor of signals

Author: sque

README.md

@@ -90,12 +90,14 @@ In order to use this feature you need to enable the celery plugin and configure
 use `current_request_id()` from inside your worker
 
 ```python
-from flask_log_request_id.extras.celery import RequestIDAwareTask
+from flask_log_request_id.extras.celery import enable_request_id_propagation
 from flask_log_request_id import current_request_id
 from celery.app import Celery
 import logging
 
-celery = Celery(task_cls=RequestIDAwareTask)
+celery = Celery()
+enable_request_id_propagation(celery)  # << This step here is critical to propagate request-id to workers
+
 app = Flask()
 
 @celery.task()

examples/server_with_celery/example_app/celery.py

@@ -1,8 +1,8 @@
 from celery import Celery, signals
-from flask_log_request_id.extras.celery import add_request_id_header, RequestIDAwareTask
+from flask_log_request_id.extras.celery import enable_request_id_propagation
 
-# Either sublcass all tasks from RequestIDAwareTask
-celery = Celery(task_cls=RequestIDAwareTask)
+celery = Celery()
+
+# You need to enable propagation on celery application
+enable_request_id_propagation(celery)
 
-# Or alternatively register the provided signal handler
-signals.before_task_publish.connect(add_request_id_header)

flask_log_request_id/extras/celery.py

@@ -1,4 +1,4 @@
-from celery import Task, current_task
+from celery import current_task, signals
 import logging as _logging
 
 from ..request_id import current_request_id
@@ -9,24 +9,26 @@
 logger = _logging.getLogger(__name__)
 
 
-class RequestIDAwareTask(Task):
+def enable_request_id_propagation(celery_app):
     """
-    Task base class that injects request id to task request object with key 'x_request_id'.
+    Will attach signal on celery application in order to propagate
+    current request id to workers
+    :param celery_app: The celery application
     """
+    signals.before_task_publish.connect(on_before_publish_insert_request_id_header)
 
-    def apply_async(self, *args, **kwargs):
-        # Set default value for 'headers' argument
-        if 'headers' not in kwargs or kwargs['headers'] is None:
-            kwargs['headers'] = {}
+
+def on_before_publish_insert_request_id_header(headers, **kwargs):
+    """
+    This function is meant to be used as signal processor for "before_task_publish".
+    :param Dict headers: The headers of the message
+    :param kwargs: Any extra keyword arguments
+    """
+    if _CELERY_X_HEADER not in headers:
         request_id = current_request_id()
+        headers[_CELERY_X_HEADER] = request_id
         logger.debug("Forwarding request_id '{}' to the task consumer.".format(request_id))
-        kwargs['headers'][_CELERY_X_HEADER] = request_id
-
-        return super(RequestIDAwareTask, self).apply_async(*args, **kwargs)
 
-def add_request_id_header(headers=None, **kwargs):
-    if _CELERY_X_HEADER not in headers:
-        headers[_CELERY_X_HEADER] = current_request_id()
 
 def ctx_celery_task_get_request_id():
     """

tests/extras/celery_tests.py

@@ -3,8 +3,7 @@
 
 from celery import Celery
 from flask_log_request_id.extras.celery import (ExecutedOutsideContext,
-                                                RequestIDAwareTask,
-                                                add_request_id_header,
+                                                on_before_publish_insert_request_id_header,
                                                 ctx_celery_task_get_request_id)
 
 
@@ -23,69 +22,27 @@ def apply_async(self, *args, **kwargs):
 class CeleryIntegrationTestCase(unittest.TestCase):
 
     @mock.patch('flask_log_request_id.extras.celery.current_request_id')
-    def test_mixin_injection(self, mocked_current_request_id):
-
-        patcher = mock.patch.object(RequestIDAwareTask, '__bases__', (MockedTask,))
-
-        with patcher:
-            patcher.is_local = True
-
-            mocked_current_request_id.return_value = 15
-            task = RequestIDAwareTask()
-            task.apply_async('test', foo='bar')
-            self.assertEqual(
-                task.apply_async_called['args'],
-                ('test', ))
-
-            self.assertDictEqual(
-                task.apply_async_called['kwargs'], {
-                    'headers': {'x_request_id': 15},
-                    'foo': 'bar'
-            })
-
-    @mock.patch('flask_log_request_id.extras.celery.current_request_id')
-    def test_issue21_called_with_headers_None(self, mocked_current_request_id):
-
-        patcher = mock.patch.object(RequestIDAwareTask, '__bases__', (MockedTask,))
-
-        with patcher:
-            patcher.is_local = True
-
-            mocked_current_request_id.return_value = 15
-            task = RequestIDAwareTask()
-            task.apply_async('test', foo='bar', headers=None)
-            self.assertEqual(
-                task.apply_async_called['args'],
-                ('test', ))
-
-            self.assertDictEqual(
-                task.apply_async_called['kwargs'], {
-                    'headers': {'x_request_id': 15},
-                    'foo': 'bar'
-            })
-
-    @mock.patch('flask_log_request_id.extras.celery.current_request_id')
-    def test_before_task_publish_hooks_adds_header(self, mocked_current_request_id):
+    def test_enable_request_id_propagation(self, mocked_current_request_id):
         mocked_current_request_id.return_value = 15
 
         headers = {}
-        add_request_id_header(headers={})
-        print(headers)
-        self.assertDictEqual(headers, {
-            'x_request_id': 15
-        })
+        on_before_publish_insert_request_id_header(headers=headers)
+        self.assertDictEqual(
+            {
+                'x_request_id': 15
+            },
+            headers)
 
     @mock.patch('flask_log_request_id.extras.celery.current_task')
     def test_ctx_fetcher_outside_context(self, mocked_current_task):
-
         mocked_current_task._get_current_object.return_value = None
         with self.assertRaises(ExecutedOutsideContext):
             ctx_celery_task_get_request_id()
 
     @mock.patch('flask_log_request_id.extras.celery.current_task')
     def test_ctx_fetcher_inside_context(self, mocked_current_task):
         mocked_current_task._get_current_object.return_value = True
-        mocked_current_task.request.get.side_effect = lambda a, default: {'x_request_id': 15, 'other':'bar'}[a]
+        mocked_current_task.request.get.side_effect = lambda a, default: {'x_request_id': 15, 'other': 'bar'}[a]
 
         self.assertEqual(ctx_celery_task_get_request_id(), 15)