feat(sqlalchemy): Support span streaming

sentry_sdk/integrations/sqlalchemy.py

@@ -1,11 +1,16 @@
 from sentry_sdk.consts import SPANSTATUS, SPANDATA
 from sentry_sdk.integrations import _check_minimum_version, Integration, DidNotEnable
-from sentry_sdk.tracing_utils import add_query_source, record_sql_queries
+from sentry_sdk.tracing_utils import (
+    add_query_source,
+    record_sql_queries_supporting_streaming,
+)
 from sentry_sdk.utils import (
     capture_internal_exceptions,
     ensure_integration_enabled,
     parse_version,
 )
+from sentry_sdk.traces import StreamedSpan, SpanStatus
+from sentry_sdk.tracing import Span
 
 try:
     from sqlalchemy.engine import Engine  # type: ignore
@@ -20,8 +25,7 @@
     from typing import Any
     from typing import ContextManager
     from typing import Optional
-
-    from sentry_sdk.tracing import Span
+    from typing import Union
 
 
 class SqlalchemyIntegration(Integration):
@@ -48,7 +52,7 @@ def _before_cursor_execute(
     executemany: bool,
     *args: "Any",
 ) -> None:
-    ctx_mgr = record_sql_queries(
+    ctx_mgr = record_sql_queries_supporting_streaming(
         cursor,
         statement,
         parameters,
@@ -78,12 +82,19 @@ def _after_cursor_execute(
         context, "_sentry_sql_span_manager", None
     )
 
+    # Record query source immediately before span is finished: accurate end timestamp and before the span is flushed.
+    span: "Optional[Union[Span, StreamedSpan]]" = getattr(
+        context, "_sentry_sql_span", None
+    )
+    if isinstance(span, StreamedSpan):
+        with capture_internal_exceptions():
+            add_query_source(span)
+
     if ctx_mgr is not None:
         context._sentry_sql_span_manager = None
         ctx_mgr.__exit__(None, None, None)
 
-    span: "Optional[Span]" = getattr(context, "_sentry_sql_span", None)
-    if span is not None:
+    if isinstance(span, Span):
         with capture_internal_exceptions():
             add_query_source(span)
 
@@ -96,7 +107,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,15 +146,20 @@ 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:
+    if isinstance(span, StreamedSpan):
+        set_on_span = span.set_attribute
+    else:
+        set_on_span = span.set_data
+
     db_system = _get_db_system(conn.engine.name)
     if db_system is not None:
-        span.set_data(SPANDATA.DB_SYSTEM, db_system)
+        set_on_span(SPANDATA.DB_SYSTEM, db_system)
 
     try:
         driver = conn.dialect.driver
         if driver:
-            span.set_data(SPANDATA.DB_DRIVER_NAME, driver)
+            set_on_span(SPANDATA.DB_DRIVER_NAME, driver)
     except Exception:
         pass
 
@@ -149,12 +168,12 @@ def _set_db_data(span: "Span", conn: "Any") -> None:
 
     db_name = conn.engine.url.database
     if db_name is not None:
-        span.set_data(SPANDATA.DB_NAME, db_name)
+        set_on_span(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)
+        set_on_span(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)
+        set_on_span(SPANDATA.SERVER_PORT, server_port)

sentry_sdk/tracing_utils.py

@@ -170,6 +170,65 @@ def record_sql_queries(
         yield span
 
 
+@contextlib.contextmanager
+def record_sql_queries_supporting_streaming(
+    cursor: "Any",
+    query: "Any",
+    params_list: "Any",
+    paramstyle: "Optional[str]",
+    executemany: bool,
+    record_cursor_repr: bool = False,
+    span_origin: str = "manual",
+) -> "Generator[Union[sentry_sdk.tracing.Span, sentry_sdk.traces.StreamedSpan], None, None]":
+    # TODO: Bring back capturing of params by default
+    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
+
+        if paramstyle == "pyformat":
+            paramstyle = "format"
+    else:
+        params_list = None
+        paramstyle = None
+
+    query = _format_sql(cursor, query)
+
+    data = {}
+    if params_list is not None:
+        data["db.params"] = params_list
+    if paramstyle is not None:
+        data["db.paramstyle"] = paramstyle
+    if executemany:
+        data["db.executemany"] = True
+    if record_cursor_repr and cursor is not None:
+        data["db.cursor"] = cursor
+
+    with capture_internal_exceptions():
+        sentry_sdk.add_breadcrumb(message=query, category="query", data=data)
+
+    if has_span_streaming_enabled(client.options):
+        with sentry_sdk.traces.start_span(
+            name="<unknown SQL query>" if query is None else 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(
     scope: "sentry_sdk.Scope", span: "sentry_sdk.tracing.Span"
 ) -> None:
@@ -316,22 +375,36 @@ def add_source(
                 span.set_attribute(SPANDATA.CODE_FUNCTION, frame.f_code.co_name)
 
 
-def add_query_source(span: "sentry_sdk.tracing.Span") -> None:
+def add_query_source(
+    span: "Union[sentry_sdk.tracing.Span, sentry_sdk.traces.StreamedSpan]",
+) -> None:
     """
     Adds OTel compatible source code information to a database query span
     """
     client = sentry_sdk.get_client()
-    if not client.is_active():
-        return
 
-    if span.timestamp is None or span.start_timestamp is None:
+    if isinstance(span, LegacySpan):
+        if not client.is_active():
+            return
+
+        # In the StreamedSpan case, we need to add the extra span information before
+        # the span finishes, so it's expected that this will be None. In the LegacySpan case,
+        # it should already be finished.
+        if span.timestamp is None:
+            return
+
+    if span.start_timestamp is None:
         return
 
     should_add_query_source = client.options.get("enable_db_query_source", True)
     if not should_add_query_source:
         return
 
-    duration = span.timestamp - span.start_timestamp
+    end_timestamp = (
+        datetime.now(timezone.utc) if span.timestamp is None else span.timestamp
+    )
+
+    duration = end_timestamp - span.start_timestamp
     threshold = client.options.get("db_query_source_threshold_ms", 0)
     slow_query = duration / timedelta(milliseconds=1) > threshold
 

tests/conftest.py

@@ -476,23 +476,34 @@ def maybe_monkeypatched_threading(request):
 
 @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:
+            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"

tests/integrations/django/test_basic.py

@@ -457,7 +457,7 @@ def test_response_trace(sentry_init, client, capture_events, render_span_tree):
 
     assert (
         '- op="view.response.render": description="serialize response"'
-        in render_span_tree(events[0])
+        in render_span_tree(events[0]["spans"], events[0]["contexts"]["trace"])
     )
 
 
@@ -596,7 +596,9 @@ def test_django_connect_trace(sentry_init, client, capture_events, render_span_t
             data = span.get("data")
             assert data.get(SPANDATA.DB_SYSTEM) == "postgresql"
 
-    assert '- op="db": description="connect"' in render_span_tree(event)
+    assert '- op="db": description="connect"' in render_span_tree(
+        event["spans"], event["contexts"]["trace"]
+    )
 
 
 @pytest.mark.forked
@@ -954,7 +956,9 @@ def test_render_spans(sentry_init, client, capture_events, render_span_tree):
         events = capture_events()
         client.get(url)
         transaction = events[0]
-        assert expected_line in render_span_tree(transaction)
+        assert expected_line in render_span_tree(
+            transaction["spans"], transaction["contexts"]["trace"]
+        )
 
 
 @pytest.mark.skipif(DJANGO_VERSION < (1, 9), reason="Requires Django >= 1.9")
@@ -1034,7 +1038,10 @@ def test_middleware_spans(sentry_init, client, capture_events, render_span_tree)
     message, transaction = events
 
     assert message["message"] == "hi"
-    assert render_span_tree(transaction) == EXPECTED_MIDDLEWARE_SPANS
+    assert (
+        render_span_tree(transaction["spans"], transaction["contexts"]["trace"])
+        == EXPECTED_MIDDLEWARE_SPANS
+    )
 
 
 def test_middleware_spans_disabled(sentry_init, client, capture_events):
@@ -1075,7 +1082,10 @@ def test_signals_spans(sentry_init, client, capture_events, render_span_tree):
     message, transaction = events
 
     assert message["message"] == "hi"
-    assert render_span_tree(transaction) == EXPECTED_SIGNALS_SPANS
+    assert (
+        render_span_tree(transaction["spans"], transaction["contexts"]["trace"])
+        == EXPECTED_SIGNALS_SPANS
+    )
 
     assert transaction["spans"][0]["op"] == "event.django"
     assert transaction["spans"][0]["description"] == "django.db.reset_queries"
@@ -1127,7 +1137,10 @@ def test_signals_spans_filtering(sentry_init, client, capture_events, render_spa
 
     (transaction,) = events
 
-    assert render_span_tree(transaction) == EXPECTED_SIGNALS_SPANS_FILTERED
+    assert (
+        render_span_tree(transaction["spans"], transaction["contexts"]["trace"])
+        == EXPECTED_SIGNALS_SPANS_FILTERED
+    )
 
     assert transaction["spans"][0]["op"] == "event.django"
     assert transaction["spans"][0]["description"] == "django.db.reset_queries"
@@ -1206,7 +1219,9 @@ def test_custom_urlconf_middleware(
     event = events.pop(0)
     assert event["transaction"] == "/custom/ok"
     if middleware_spans:
-        assert "custom_urlconf_middleware" in render_span_tree(event)
+        assert "custom_urlconf_middleware" in render_span_tree(
+            event["spans"], event["contexts"]["trace"]
+        )
 
     _content, status, _headers = unpack_werkzeug_response(client.get("/custom/exc"))
     assert status.lower() == "500 internal server error"
@@ -1216,7 +1231,9 @@ def test_custom_urlconf_middleware(
     assert error_event["exception"]["values"][-1]["mechanism"]["type"] == "django"
     assert transaction_event["transaction"] == "/custom/exc"
     if middleware_spans:
-        assert "custom_urlconf_middleware" in render_span_tree(transaction_event)
+        assert "custom_urlconf_middleware" in render_span_tree(
+            transaction_event["spans"], transaction_event["contexts"]["trace"]
+        )
 
     settings.MIDDLEWARE.pop(0)