Enable request tracing on Celery success/fail logs

Previously these logs wouldn't have a Request ID attached since the
Celery hooks run after the __call__ method where we enable request
tracing for normal application logs. For the failure log especially
it will be useful to have this feature.

app/celery/celery.py

@@ -1,4 +1,5 @@
 import time
+from contextlib import contextmanager
 
 from celery import Celery, Task
 from celery.signals import worker_process_shutdown
@@ -35,7 +36,16 @@ def make_task(app):
             # task context (aka "request").
             return self.request.get('notify_request_id')
 
+        @contextmanager
+        def app_context(self):
+            with app.app_context():
+                # Add 'request_id' to 'g' so that it gets logged.
+                g.request_id = self.request_id
+                yield
+
         def on_success(self, retval, task_id, args, kwargs):
+            # enables request id tracing for these logs
+            with self.app_context():
                 elapsed_time = time.monotonic() - self.start
 
                 app.logger.info(
@@ -54,6 +64,8 @@ def make_task(app):
                 )
 
         def on_failure(self, exc, task_id, args, kwargs, einfo):
+            # enables request id tracing for these logs
+            with self.app_context():
                 app.logger.exception(
                     "Celery task {task_name} (queue: {queue_name}) failed".format(
                         task_name=self.name,
@@ -72,12 +84,10 @@ def make_task(app):
 
         def __call__(self, *args, **kwargs):
             # ensure task has flask context to access config, logger, etc
-            with app.app_context():
+            with self.app_context():
                 self.start = time.monotonic()
                 # TEMPORARY: remove old piggyback values from kwargs
                 kwargs.pop('request_id', None)
-                # Add 'request_id' to 'g' so that it gets logged. Note
-                g.request_id = self.request_id
 
                 return super().__call__(*args, **kwargs)