feat(django): Support span streaming

sentry_sdk/integrations/asyncpg.py

@@ -12,7 +12,7 @@
 from sentry_sdk.tracing_utils import (
     add_query_source,
     has_span_streaming_enabled,
-    record_sql_queries_supporting_streaming,
+    record_sql_queries,
 )
 from sentry_sdk.utils import (
     capture_internal_exceptions,
@@ -82,7 +82,7 @@ async def _inner(*args: "Any", **kwargs: "Any") -> "T":
             return await f(*args, **kwargs)
 
         query = _normalize_query(args[1])
-        with record_sql_queries_supporting_streaming(
+        with record_sql_queries(
             cursor=None,
             query=query,
             params_list=None,
@@ -123,7 +123,7 @@ def _record(
     param_style = "pyformat" if params_list else None
 
     query = _normalize_query(query)
-    with record_sql_queries_supporting_streaming(
+    with record_sql_queries(
         cursor=cursor,
         query=query,
         params_list=params_list,
@@ -170,7 +170,7 @@ async def _inner(*args: "Any", **kwargs: "Any") -> "T":
 
         cursor = args[0]
         query = _normalize_query(cursor._query)
-        with record_sql_queries_supporting_streaming(
+        with record_sql_queries(
             cursor=cursor,
             query=query,
             params_list=None,

sentry_sdk/integrations/django/__init__.py

@@ -15,8 +15,13 @@
 from sentry_sdk.integrations.wsgi import SentryWsgiMiddleware
 from sentry_sdk.scope import add_global_event_processor, should_send_default_pii
 from sentry_sdk.serializer import add_global_repr_processor, add_repr_sequence_type
+from sentry_sdk.traces import StreamedSpan
 from sentry_sdk.tracing import SOURCE_FOR_STYLE, TransactionSource
-from sentry_sdk.tracing_utils import add_query_source, record_sql_queries
+from sentry_sdk.tracing_utils import (
+    add_query_source,
+    has_span_streaming_enabled,
+    record_sql_queries,
+)
 from sentry_sdk.utils import (
     CONTEXTVARS_ERROR_MESSAGE,
     HAS_REAL_CONTEXTVARS,
@@ -82,6 +87,7 @@
 
     from sentry_sdk._types import Event, EventProcessor, Hint, NotImplementedType
     from sentry_sdk.integrations.wsgi import _ScopedResponse
+    from sentry_sdk.traces import StreamedSpan
     from sentry_sdk.tracing import Span
 
 
@@ -639,8 +645,13 @@ def execute(
             _set_db_data(span, self)
             result = real_execute(self, sql, params)
 
-        with capture_internal_exceptions():
-            add_query_source(span)
+            if isinstance(span, StreamedSpan):
+                with capture_internal_exceptions():
+                    add_query_source(span)
+
+        if not isinstance(span, StreamedSpan):
+            with capture_internal_exceptions():
+                add_query_source(span)
 
         return result
 
@@ -660,8 +671,13 @@ def executemany(
 
             result = real_executemany(self, sql, param_list)
 
-        with capture_internal_exceptions():
-            add_query_source(span)
+            if isinstance(span, StreamedSpan):
+                with capture_internal_exceptions():
+                    add_query_source(span)
+
+        if not isinstance(span, StreamedSpan):
+            with capture_internal_exceptions():
+                add_query_source(span)
 
         return result
 
@@ -670,41 +686,77 @@ def connect(self: "BaseDatabaseWrapper") -> None:
         with capture_internal_exceptions():
             sentry_sdk.add_breadcrumb(message="connect", category="query")
 
-        with sentry_sdk.start_span(
-            op=OP.DB,
-            name="connect",
-            origin=DjangoIntegration.origin_db,
-        ) as span:
-            _set_db_data(span, self)
-            return real_connect(self)
+        span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
+        if span_streaming:
+            with sentry_sdk.traces.start_span(
+                name="connect",
+                attributes={
+                    "sentry.op": OP.DB,
+                    "sentry.origin": DjangoIntegration.origin_db,
+                },
+            ) as span:
+                _set_db_data(span, self)
+                return real_connect(self)
+        else:
+            with sentry_sdk.start_span(
+                op=OP.DB,
+                name="connect",
+                origin=DjangoIntegration.origin_db,
+            ) as span:
+                _set_db_data(span, self)
+                return real_connect(self)
 
     def _commit(self: "BaseDatabaseWrapper") -> None:
         integration = sentry_sdk.get_client().get_integration(DjangoIntegration)
 
         if integration is None or not integration.db_transaction_spans:
             return real_commit(self)
 
-        with sentry_sdk.start_span(
-            op=OP.DB,
-            name=SPANNAME.DB_COMMIT,
-            origin=DjangoIntegration.origin_db,
-        ) as span:
-            _set_db_data(span, self, SPANNAME.DB_COMMIT)
-            return real_commit(self)
+        span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
+        if span_streaming:
+            with sentry_sdk.traces.start_span(
+                name=SPANNAME.DB_COMMIT,
+                attributes={
+                    "sentry.op": OP.DB,
+                    "sentry.origin": DjangoIntegration.origin_db,
+                },
+            ) as span:
+                _set_db_data(span, self, SPANNAME.DB_COMMIT)
+                return real_commit(self)
+        else:
+            with sentry_sdk.start_span(
+                op=OP.DB,
+                name=SPANNAME.DB_COMMIT,
+                origin=DjangoIntegration.origin_db,
+            ) as span:
+                _set_db_data(span, self, SPANNAME.DB_COMMIT)
+                return real_commit(self)
 
     def _rollback(self: "BaseDatabaseWrapper") -> None:
         integration = sentry_sdk.get_client().get_integration(DjangoIntegration)
 
         if integration is None or not integration.db_transaction_spans:
             return real_rollback(self)
 
-        with sentry_sdk.start_span(
-            op=OP.DB,
-            name=SPANNAME.DB_ROLLBACK,
-            origin=DjangoIntegration.origin_db,
-        ) as span:
-            _set_db_data(span, self, SPANNAME.DB_ROLLBACK)
-            return real_rollback(self)
+        span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
+        if span_streaming:
+            with sentry_sdk.traces.start_span(
+                name=SPANNAME.DB_ROLLBACK,
+                attributes={
+                    "sentry.op": OP.DB,
+                    "sentry.origin": DjangoIntegration.origin_db,
+                },
+            ) as span:
+                _set_db_data(span, self, SPANNAME.DB_ROLLBACK)
+                return real_rollback(self)
+        else:
+            with sentry_sdk.start_span(
+                op=OP.DB,
+                name=SPANNAME.DB_ROLLBACK,
+                origin=DjangoIntegration.origin_db,
+            ) as span:
+                _set_db_data(span, self, SPANNAME.DB_ROLLBACK)
+                return real_rollback(self)
 
     CursorWrapper.execute = execute
     CursorWrapper.executemany = executemany
@@ -715,14 +767,22 @@ def _rollback(self: "BaseDatabaseWrapper") -> None:
 
 
 def _set_db_data(
-    span: "Span", cursor_or_db: "Any", db_operation: "Optional[str]" = None
+    span: "Union[Span, StreamedSpan]",
+    cursor_or_db: "Any",
+    db_operation: "Optional[str]" = None,
 ) -> None:
     db = cursor_or_db.db if hasattr(cursor_or_db, "db") else cursor_or_db
     vendor = db.vendor
-    span.set_data(SPANDATA.DB_SYSTEM, vendor)
+    if isinstance(span, StreamedSpan):
+        span.set_attribute(SPANDATA.DB_SYSTEM_NAME, vendor)
 
-    if db_operation is not None:
-        span.set_data(SPANDATA.DB_OPERATION, db_operation)
+        if db_operation is not None:
+            span.set_attribute(SPANDATA.DB_OPERATION_NAME, db_operation)
+    else:
+        span.set_data(SPANDATA.DB_SYSTEM, vendor)
+
+        if db_operation is not None:
+            span.set_data(SPANDATA.DB_OPERATION, db_operation)
 
     # Some custom backends override `__getattr__`, making it look like `cursor_or_db`
     # actually has a `connection` and the `connection` has a `get_dsn_parameters`
@@ -754,20 +814,29 @@ def _set_db_data(
             connection_params = db.get_connection_params()
 
     db_name = connection_params.get("dbname") or connection_params.get("database")
-    if db_name is not None:
-        span.set_data(SPANDATA.DB_NAME, db_name)
+
+    if isinstance(span, StreamedSpan):
+        if db_name is not None:
+            span.set_attribute(SPANDATA.DB_NAMESPACE, db_name)
+
+        set_on_span = span.set_attribute
+    else:
+        if db_name is not None:
+            span.set_data(SPANDATA.DB_NAME, db_name)
+
+        set_on_span = span.set_data
 
     server_address = connection_params.get("host")
     if server_address is not None:
-        span.set_data(SPANDATA.SERVER_ADDRESS, server_address)
+        set_on_span(SPANDATA.SERVER_ADDRESS, server_address)
 
     server_port = connection_params.get("port")
     if server_port is not None:
-        span.set_data(SPANDATA.SERVER_PORT, str(server_port))
+        set_on_span(SPANDATA.SERVER_PORT, str(server_port))
 
     server_socket_address = connection_params.get("unix_socket")
     if server_socket_address is not None:
-        span.set_data(SPANDATA.SERVER_SOCKET_ADDRESS, server_socket_address)
+        set_on_span(SPANDATA.SERVER_SOCKET_ADDRESS, server_socket_address)
 
 
 def add_template_context_repr_sequence() -> None:

sentry_sdk/integrations/django/asgi.py

@@ -17,6 +17,8 @@
 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,
     ensure_integration_enabled,
@@ -168,24 +170,42 @@ 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)
 
-        with sentry_sdk.start_span(
-            op=OP.VIEW_RENDER,
-            name=request.resolver_match.view_name,
-            origin=DjangoIntegration.origin,
-        ):
-            return await callback(request, *args, **kwargs)
+        if span_streaming:
+            with sentry_sdk.traces.start_span(
+                name=request.resolver_match.view_name,
+                attributes={
+                    "sentry.op": OP.VIEW_RENDER,
+                    "sentry.origin": DjangoIntegration.origin,
+                },
+            ):
+                return await callback(request, *args, **kwargs)
+        else:
+            with sentry_sdk.start_span(
+                op=OP.VIEW_RENDER,
+                name=request.resolver_match.view_name,
+                origin=DjangoIntegration.origin,
+            ):
+                return await callback(request, *args, **kwargs)
 
     return sentry_wrapped_callback