feat(sqlalchemy): Support span streaming

sentry_sdk/integrations/sqlalchemy.py

@@ -6,6 +6,7 @@
     ensure_integration_enabled,
     parse_version,
 )
+from sentry_sdk.traces import StreamedSpan, SpanStatus
 
 try:
     from sqlalchemy.engine import Engine  # type: ignore
@@ -20,6 +21,7 @@
     from typing import Any
     from typing import ContextManager
     from typing import Optional
+    from typing import Union
 
     from sentry_sdk.tracing import Span
 
@@ -96,7 +98,10 @@ def _handle_error(context: "Any", *args: "Any") -> None:
     span: "Optional[Span]" = getattr(execution_context, "_sentry_sql_span", None)
 
     if span is not None:
-        span.set_status(SPANSTATUS.INTERNAL_ERROR)
+        if isinstance(span, StreamedSpan):
+            span.status = SpanStatus.ERROR
+        else:
+            span.set_status(SPANSTATUS.INTERNAL_ERROR)
 
     # _after_cursor_execute does not get called for crashing SQL stmts. Judging
     # from SQLAlchemy codebase it does seem like any error coming into this
@@ -132,29 +137,53 @@ def _get_db_system(name: str) -> "Optional[str]":
     return None
 
 
-def _set_db_data(span: "Span", conn: "Any") -> None:
+def _set_db_data(span: "Union[Span, StreamedSpan]", conn: "Any") -> None:
     db_system = _get_db_system(conn.engine.name)
-    if db_system is not None:
-        span.set_data(SPANDATA.DB_SYSTEM, db_system)
-
-    try:
-        driver = conn.dialect.driver
-        if driver:
-            span.set_data(SPANDATA.DB_DRIVER_NAME, driver)
-    except Exception:
-        pass
+    if isinstance(span, StreamedSpan):
+        if db_system is not None:
+            span.set_attribute(SPANDATA.DB_SYSTEM, db_system)
+
+        try:
+            driver = conn.dialect.driver
+            if driver:
+                span.set_attribute(SPANDATA.DB_DRIVER_NAME, driver)
+        except Exception:
+            pass
+    else:
+        if db_system is not None:
+            span.set_data(SPANDATA.DB_SYSTEM, db_system)
+
+        try:
+            driver = conn.dialect.driver
+            if driver:
+                span.set_data(SPANDATA.DB_DRIVER_NAME, driver)
+        except Exception:
+            pass
 
     if conn.engine.url is None:
         return
 
-    db_name = conn.engine.url.database
-    if db_name is not None:
-        span.set_data(SPANDATA.DB_NAME, db_name)
-
-    server_address = conn.engine.url.host
-    if server_address is not None:
-        span.set_data(SPANDATA.SERVER_ADDRESS, server_address)
-
-    server_port = conn.engine.url.port
-    if server_port is not None:
-        span.set_data(SPANDATA.SERVER_PORT, server_port)
+    if isinstance(span, StreamedSpan):
+        db_name = conn.engine.url.database
+        if db_name is not None:
+            span.set_attribute(SPANDATA.DB_NAME, db_name)
+
+        server_address = conn.engine.url.host
+        if server_address is not None:
+            span.set_attribute(SPANDATA.SERVER_ADDRESS, server_address)
+
+        server_port = conn.engine.url.port
+        if server_port is not None:
+            span.set_attribute(SPANDATA.SERVER_PORT, server_port)
+    else:
+        db_name = conn.engine.url.database
+        if db_name is not None:
+            span.set_data(SPANDATA.DB_NAME, db_name)
+
+        server_address = conn.engine.url.host
+        if server_address is not None:
+            span.set_data(SPANDATA.SERVER_ADDRESS, server_address)
+
+        server_port = conn.engine.url.port
+        if server_port is not None:
+            span.set_data(SPANDATA.SERVER_PORT, server_port)

sentry_sdk/tracing_utils.py

@@ -135,7 +135,8 @@
     span_origin: str = "manual",
 ) -> "Generator[sentry_sdk.tracing.Span, None, None]":
     # TODO: Bring back capturing of params by default
-    if sentry_sdk.get_client().options["_experiments"].get("record_sql_params", False):
+    client = sentry_sdk.get_client()
+    if client.options["_experiments"].get("record_sql_params", False):
         if not params_list or params_list == [None]:
             params_list = None
 
@@ -160,14 +161,26 @@
     with capture_internal_exceptions():
         sentry_sdk.add_breadcrumb(message=query, category="query", data=data)
 
-    with sentry_sdk.start_span(
-        op=OP.DB,
-        name=query,
-        origin=span_origin,
-    ) as span:
-        for k, v in data.items():
-            span.set_data(k, v)
-        yield span
+    if has_span_streaming_enabled(client.options):
+        with sentry_sdk.traces.start_span(
+            name=query,
+            attributes={
+                "sentry.origin": span_origin,
+                "sentry.op": OP.DB,
+            },
+        ) as span:
+            for k, v in data.items():
+                span.set_attribute(k, v)
+            yield span
+    else:
+        with sentry_sdk.start_span(
+            op=OP.DB,
+            name=query,
+            origin=span_origin,
+        ) as span:
+            for k, v in data.items():
+                span.set_data(k, v)
+            yield span
 
 
 def maybe_create_breadcrumbs_from_span(

tests/conftest.py

@@ -476,23 +476,35 @@
 
 @pytest.fixture
 def render_span_tree():
-    def inner(event):
-        assert event["type"] == "transaction"
+    def inner(spans, root_span=None):
+        streamed_spans = False
+        if root_span is None:
+            streamed_spans = True
 
         by_parent = {}
-        for span in event["spans"]:
+        for span in spans:
+            print(span)
+            if "parent_span_id" not in span:
+                root_span = span
+                continue
+
             by_parent.setdefault(span["parent_span_id"], []).append(span)
 
         def render_span(span):
-            yield "- op={}: description={}".format(
-                json.dumps(span.get("op")), json.dumps(span.get("description"))
-            )
+            if streamed_spans:
+                yield "- sentry.op={}: name={}".format(
+                    json.dumps(span["attributes"].get("sentry.op")),
+                    json.dumps(span["name"]),
+                )
+            else:
+                yield "- op={}: description={}".format(
+                    json.dumps(span.get("op")), json.dumps(span.get("description"))
+                )
+
             for subspan in by_parent.get(span["span_id"]) or ():
                 for line in render_span(subspan):
                     yield "  {}".format(line)
 
-        root_span = event["contexts"]["trace"]
-
         return "\n".join(render_span(root_span))
 
     return inner

tests/integrations/django/asgi/test_asgi.py

@@ -292,8 +292,9 @@ async def test_async_middleware_spans(
 
     (transaction,) = events
 
+    assert transaction["type"] == "transaction"
     assert (
-        render_span_tree(transaction)
+        render_span_tree(transaction["spans"], transaction["contexts"]["trace"])
         == """\
 - op="http.server": description=null
   - op="event.django": description="django.db.reset_queries"