From ceb7cee0096933ff96355049a614052d70a92f8d Mon Sep 17 00:00:00 2001 From: Katie Smith Date: Wed, 23 Oct 2019 16:14:59 +0100 Subject: [PATCH] Pass request_id to tasks if available We want to pass the `request_id` to Celery tasks if the task is called from an HTTP request, so that we can add the `request_id` to the logs. This change overwrites `apply_async` to add the `request_id` to the kwargs if available. When we call the task, we then add the `request_id` to g on Flask's application context. Tasks called from `send_task` won't have a `request_id` for now, and this change only affects tasks called from HTTP requests (not from other tasks or from Celery Beat). --- app/celery/celery.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/app/celery/celery.py b/app/celery/celery.py index 56c7f72b2..506de84d8 100644 --- a/app/celery/celery.py +++ b/app/celery/celery.py @@ -2,7 +2,8 @@ import time from celery import Celery, Task from celery.signals import worker_process_shutdown -from flask import current_app +from flask import current_app, g, request +from flask.ctx import has_request_context @worker_process_shutdown.connect @@ -32,8 +33,20 @@ def make_task(app): # ensure task has flask context to access config, logger, etc with app.app_context(): self.start = time.time() + # Remove 'request_id' from the kwargs (so the task doesn't get an unexpected kwarg), then add it to g + # so that it gets logged + g.request_id = kwargs.pop('request_id', None) return super().__call__(*args, **kwargs) + def apply_async(self, args=None, kwargs=None, task_id=None, producer=None, + link=None, link_error=None, **options): + kwargs = kwargs or {} + + if has_request_context() and hasattr(request, 'request_id'): + kwargs['request_id'] = request.request_id + + return super().apply_async(args, kwargs, task_id, producer, link, link_error, **options) + return NotifyTask