From 36666c81c676a46a711648f4a6b3392b1d8a126a Mon Sep 17 00:00:00 2001 From: Jimmy Royer Date: Wed, 21 Aug 2024 13:56:12 -0400 Subject: [PATCH] Wire the beat signals with xray (#2259) --- app/celery/celery.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/app/celery/celery.py b/app/celery/celery.py index be5a5593aa..28cd2d3153 100644 --- a/app/celery/celery.py +++ b/app/celery/celery.py @@ -2,7 +2,14 @@ from flask import current_app -from celery import Celery, Task +from app.aws.xray_celery_handlers import ( + xray_after_task_publish, + xray_before_task_publish, + xray_task_failure, + xray_task_postrun, + xray_task_prerun, +) +from celery import Celery, Task, signals from celery.signals import worker_process_shutdown @@ -42,6 +49,14 @@ def init_app(self, app): task_cls=make_task(app), ) + # Register the xray handlers + signals.after_task_publish.connect(xray_after_task_publish) + signals.before_task_publish.connect(xray_before_task_publish) + signals.task_failure.connect(xray_task_failure) + signals.task_postrun.connect(xray_task_postrun) + signals.task_prerun.connect(xray_task_prerun) + signals.beat_init.connect(xray_task_prerun) + # See https://docs.celeryproject.org/en/stable/userguide/configuration.html self.conf.update( {