set thread id

Author: alexander-alderman-webb

sentry_sdk/integrations/django/asgi.py

@@ -17,6 +17,7 @@
 from sentry_sdk.consts import OP
 from sentry_sdk.integrations.asgi import SentryAsgiMiddleware
 from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.traces import StreamedSpan
 from sentry_sdk.tracing_utils import has_span_streaming_enabled
 from sentry_sdk.utils import (
     capture_internal_exceptions,
@@ -169,19 +170,26 @@ def wrap_async_view(callback: "Any") -> "Any":
     async def sentry_wrapped_callback(
         request: "Any", *args: "Any", **kwargs: "Any"
     ) -> "Any":
+        client = sentry_sdk.get_client()
+        span_streaming = has_span_streaming_enabled(client.options)
         current_scope = sentry_sdk.get_current_scope()
-        if current_scope.transaction is not None:
-            current_scope.transaction.update_active_thread()
+        if span_streaming:
+            current_span = current_scope.streamed_span
+            if type(current_span) is StreamedSpan:
+                segment = current_span._segment
+                segment._update_active_thread()
+        else:
+            if current_scope.transaction is not None:
+                current_scope.transaction.update_active_thread()
 
         sentry_scope = sentry_sdk.get_isolation_scope()
         if sentry_scope.profile is not None:
             sentry_scope.profile.update_active_thread_id()
 
-        integration = sentry_sdk.get_client().get_integration(DjangoIntegration)
+        integration = client.get_integration(DjangoIntegration)
         if not integration or not integration.middleware_spans:
             return await callback(request, *args, **kwargs)
 
-        span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
         if span_streaming:
             with sentry_sdk.traces.start_span(
                 name=request.resolver_match.view_name,

sentry_sdk/integrations/django/views.py

@@ -3,6 +3,7 @@
 
 import sentry_sdk
 from sentry_sdk.consts import OP
+from sentry_sdk.traces import StreamedSpan
 from sentry_sdk.tracing_utils import has_span_streaming_enabled
 
 if TYPE_CHECKING:
@@ -84,21 +85,28 @@ def _wrap_sync_view(callback: "Any") -> "Any":
 
     @functools.wraps(callback)
     def sentry_wrapped_callback(request: "Any", *args: "Any", **kwargs: "Any") -> "Any":
+        client = sentry_sdk.get_client()
+        span_streaming = has_span_streaming_enabled(client.options)
         current_scope = sentry_sdk.get_current_scope()
-        if current_scope.transaction is not None:
-            current_scope.transaction.update_active_thread()
+        if span_streaming:
+            current_span = current_scope.streamed_span
+            if type(current_span) is StreamedSpan:
+                segment = current_span._segment
+                segment._update_active_thread()
+        else:
+            if current_scope.transaction is not None:
+                current_scope.transaction.update_active_thread()
 
         sentry_scope = sentry_sdk.get_isolation_scope()
         # set the active thread id to the handler thread for sync views
         # this isn't necessary for async views since that runs on main
         if sentry_scope.profile is not None:
             sentry_scope.profile.update_active_thread_id()
 
-        integration = sentry_sdk.get_client().get_integration(DjangoIntegration)
+        integration = client.get_integration(DjangoIntegration)
         if not integration or not integration.middleware_spans:
             return callback(request, *args, **kwargs)
 
-        span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
         if span_streaming:
             with sentry_sdk.traces.start_span(
                 name=request.resolver_match.view_name,

tests/integrations/django/asgi/test_asgi.py

@@ -228,11 +228,13 @@ async def test_active_thread_id(
 
         data = json.loads(response["body"])
 
-        spans = [item for item in items if item.type == "span"]
+        sentry_sdk.flush()
+        spans = [item.payload for item in items if item.type == "span"]
 
         for span in spans:
-            trace_context = span["contexts"]["trace"]
-            assert str(data["active"]) == trace_context["attributes"]["thread.id"]
+            if span["is_segment"] is False:
+                continue
+            assert str(data["active"]) == span["attributes"]["thread.id"]
     else:
         with mock.patch(
             "sentry_sdk.profiler.transaction_profiler.PROFILE_MINIMUM_SAMPLES", 0