adding on/off switch for xray (#2261)

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(

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", "{}"))

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)
 

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()