From 70c478fc0cf3fcec7c73776a8e849acaae5d2d87 Mon Sep 17 00:00:00 2001
From: Mike Pond <32133001+P0NDER0SA@users.noreply.github.com>
Date: Wed, 21 Aug 2024 14:40:20 -0400
Subject: [PATCH] adding on/off switch for xray (#2261)

---
 app/celery/celery.py | 15 ++++++++-------
 app/config.py        |  3 +++
 application.py       |  5 +++--
 run_celery.py        |  5 +++--
 4 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/app/celery/celery.py b/app/celery/celery.py
index 28cd2d3153..31b0297618 100644
--- a/app/celery/celery.py
+++ b/app/celery/celery.py
@@ -49,13 +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)
+        if app.config["AWS_XRAY_ENABLED"] == "true":
+            # 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(
diff --git a/app/config.py b/app/config.py
index fcb6aa4297..9d2885dfb9 100644
--- a/app/config.py
+++ b/app/config.py
@@ -240,6 +240,9 @@ class Config(object):
     DEBUG = False
     NOTIFY_LOG_PATH = os.getenv("NOTIFY_LOG_PATH")
 
+    # AWS Xray
+    AWS_XRAY_ENABLED = os.getenv("AWS_XRAY_ENABLED", "false")
+
     # Cronitor
     CRONITOR_ENABLED = False
     CRONITOR_KEYS = json.loads(os.getenv("CRONITOR_KEYS", "{}"))
diff --git a/application.py b/application.py
index 12cea1703c..07553cb347 100644
--- a/application.py
+++ b/application.py
@@ -20,8 +20,9 @@
 
 app = create_app(application)
 
-xray_recorder.configure(service='api')
-XRayMiddleware(app, xray_recorder)
+if app.config["AWS_XRAY_ENABLED"] == "true":
+    xray_recorder.configure(service='api')
+    XRayMiddleware(app, xray_recorder)
 
 apig_wsgi_handler = make_lambda_handler(app, binary_support=True)
 
diff --git a/run_celery.py b/run_celery.py
index 3a5cf9aa95..635b5653c8 100644
--- a/run_celery.py
+++ b/run_celery.py
@@ -15,7 +15,8 @@
 application = Flask("celery")
 create_app(application)
 
-xray_recorder.configure(service='celery')
-XRayMiddleware(application, xray_recorder)
+if application.config["AWS_XRAY_ENABLED"] == "true":
+    xray_recorder.configure(service='celery')
+    XRayMiddleware(application, xray_recorder)
 
 application.app_context().push()