From 5520bdb05eacc69c78713557dae5788549ece052 Mon Sep 17 00:00:00 2001
From: Anton Pirker <anton.pirker@sentry.io>
Date: Mon, 6 May 2024 13:32:43 +0200
Subject: [PATCH] Fix DSC in celery tasks started by Celery Beat. (#3047)

Only create the span for enqueuing the task when we are not currently running in a celery beat task. (because then there is only the span without a transaction and thus the baggage header can not be given to the child celery task.) Without the span the child celery task creates its own trace, this is what we want.
---
 .../lambda_function.py                        | 18 ++++++-------
 sentry_sdk/integrations/celery/__init__.py    | 25 ++++++++++++++++---
 sentry_sdk/tracing_utils.py                   | 10 ++++++++
 3 files changed, 41 insertions(+), 12 deletions(-)

diff --git a/scripts/aws_lambda_functions/sentryPythonDeleteTestFunctions/lambda_function.py b/scripts/aws_lambda_functions/sentryPythonDeleteTestFunctions/lambda_function.py
index 1fc3994176..ce7afb6aa4 100644
--- a/scripts/aws_lambda_functions/sentryPythonDeleteTestFunctions/lambda_function.py
+++ b/scripts/aws_lambda_functions/sentryPythonDeleteTestFunctions/lambda_function.py
@@ -1,12 +1,12 @@
 import boto3
-import sentry_sdk 
+import sentry_sdk
 
 
 monitor_slug = "python-sdk-aws-lambda-tests-cleanup"
 monitor_config = {
     "schedule": {
         "type": "crontab",
-        "value": "0 12 * * 0", # 12 o'clock on Sunday
+        "value": "0 12 * * 0",  # 12 o'clock on Sunday
     },
     "timezone": "UTC",
     "checkin_margin": 2,
@@ -24,7 +24,7 @@ def delete_lambda_functions(prefix="test_"):
     """
     client = boto3.client("lambda", region_name="us-east-1")
     functions_deleted = 0
-    
+
     functions_paginator = client.get_paginator("list_functions")
     for functions_page in functions_paginator.paginate():
         for func in functions_page["Functions"]:
@@ -39,17 +39,17 @@ def delete_lambda_functions(prefix="test_"):
                     print(f"Got exception: {ex}")
 
     return functions_deleted
-    
+
 
 def lambda_handler(event, context):
     functions_deleted = delete_lambda_functions()
-    
+
     sentry_sdk.metrics.gauge(
-        key="num_aws_functions_deleted", 
+        key="num_aws_functions_deleted",
         value=functions_deleted,
     )
-    
+
     return {
-        'statusCode': 200,
-        'body': f"{functions_deleted} AWS Lambda functions deleted successfully."
+        "statusCode": 200,
+        "body": f"{functions_deleted} AWS Lambda functions deleted successfully.",
     }
diff --git a/sentry_sdk/integrations/celery/__init__.py b/sentry_sdk/integrations/celery/__init__.py
index 74205a0184..62fdb1da6f 100644
--- a/sentry_sdk/integrations/celery/__init__.py
+++ b/sentry_sdk/integrations/celery/__init__.py
@@ -30,6 +30,7 @@
     from typing import List
     from typing import Optional
     from typing import TypeVar
+    from typing import Union
 
     from sentry_sdk._types import EventProcessor, Event, Hint, ExcInfo
     from sentry_sdk.tracing import Span
@@ -223,6 +224,16 @@ def _update_celery_task_headers(original_headers, span, monitor_beat_tasks):
     return updated_headers
 
 
+class NoOpMgr:
+    def __enter__(self):
+        # type: () -> None
+        return None
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        # type: (Any, Any, Any) -> None
+        return None
+
+
 def _wrap_apply_async(f):
     # type: (F) -> F
     @wraps(f)
@@ -242,9 +253,17 @@ def apply_async(*args, **kwargs):
 
         task = args[0]
 
-        with sentry_sdk.start_span(
-            op=OP.QUEUE_SUBMIT_CELERY, description=task.name
-        ) as span:
+        task_started_from_beat = (
+            sentry_sdk.Scope.get_isolation_scope()._name == "celery-beat"
+        )
+
+        span_mgr = (
+            sentry_sdk.start_span(op=OP.QUEUE_SUBMIT_CELERY, description=task.name)
+            if not task_started_from_beat
+            else NoOpMgr()
+        )  # type: Union[Span, NoOpMgr]
+
+        with span_mgr as span:
             kwargs["headers"] = _update_celery_task_headers(
                 kwarg_headers, span, integration.monitor_beat_tasks
             )
diff --git a/sentry_sdk/tracing_utils.py b/sentry_sdk/tracing_utils.py
index 556a466c0b..fac51f4848 100644
--- a/sentry_sdk/tracing_utils.py
+++ b/sentry_sdk/tracing_utils.py
@@ -421,6 +421,16 @@ def update(self, other_dict):
             except AttributeError:
                 pass
 
+    def __repr__(self):
+        # type: (...) -> str
+        return "<PropagationContext _trace_id={} _span_id={} parent_span_id={} parent_sampled={} dynamic_sampling_context={}>".format(
+            self._trace_id,
+            self._span_id,
+            self.parent_span_id,
+            self.parent_sampled,
+            self.dynamic_sampling_context,
+        )
+
 
 class Baggage:
     """