use separate path for streaming

Author: alexander-alderman-webb

sentry_sdk/integrations/sqlalchemy.py

@@ -1,6 +1,9 @@
 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,
@@ -49,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,

sentry_sdk/tracing_utils.py

@@ -133,6 +133,52 @@ def record_sql_queries(
     executemany: bool,
     record_cursor_repr: bool = False,
     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):
+        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)
+
+    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
+
+
+@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()

tests/integrations/sqlalchemy/test_sqlalchemy.py

@@ -14,7 +14,7 @@
 from sentry_sdk.consts import DEFAULT_MAX_VALUE_LENGTH, SPANDATA
 from sentry_sdk.integrations.sqlalchemy import SqlalchemyIntegration
 from sentry_sdk.serializer import MAX_EVENT_BYTES
-from sentry_sdk.tracing_utils import record_sql_queries
+from sentry_sdk.tracing_utils import record_sql_queries_supporting_streaming
 from sentry_sdk.utils import json_dumps
 
 
@@ -922,7 +922,9 @@ class Person(Base):
 
             class fake_record_sql_queries:  # noqa: N801
                 def __init__(self, *args, **kwargs):
-                    with record_sql_queries(*args, **kwargs) as span:
+                    with record_sql_queries_supporting_streaming(
+                        *args, **kwargs
+                    ) as span:
                         self.span = span
 
                     if span_streaming:
@@ -939,7 +941,7 @@ def __exit__(self, type, value, traceback):
                     pass
 
             with mock.patch(
-                "sentry_sdk.integrations.sqlalchemy.record_sql_queries",
+                "sentry_sdk.integrations.sqlalchemy.record_sql_queries_supporting_streaming",
                 fake_record_sql_queries,
             ):
                 assert session.query(Person).first() == bob
@@ -983,7 +985,9 @@ class Person(Base):
 
             class fake_record_sql_queries:  # noqa: N801
                 def __init__(self, *args, **kwargs):
-                    with record_sql_queries(*args, **kwargs) as span:
+                    with record_sql_queries_supporting_streaming(
+                        *args, **kwargs
+                    ) as span:
                         self.span = span
 
                     if span_streaming:
@@ -1000,7 +1004,7 @@ def __exit__(self, type, value, traceback):
                     pass
 
             with mock.patch(
-                "sentry_sdk.integrations.sqlalchemy.record_sql_queries",
+                "sentry_sdk.integrations.sqlalchemy.record_sql_queries_supporting_streaming",
                 fake_record_sql_queries,
             ):
                 assert session.query(Person).first() == bob
@@ -1061,7 +1065,9 @@ class Person(Base):
 
             class fake_record_sql_queries:  # noqa: N801
                 def __init__(self, *args, **kwargs):
-                    with record_sql_queries(*args, **kwargs) as span:
+                    with record_sql_queries_supporting_streaming(
+                        *args, **kwargs
+                    ) as span:
                         self.span = span
 
                     self.span._start_timestamp = datetime(2024, 1, 1, microsecond=0)
@@ -1074,7 +1080,7 @@ def __exit__(self, type, value, traceback):
                     pass
 
             with mock.patch(
-                "sentry_sdk.integrations.sqlalchemy.record_sql_queries",
+                "sentry_sdk.integrations.sqlalchemy.record_sql_queries_supporting_streaming",
                 fake_record_sql_queries,
             ):
                 assert session.query(Person).first() == bob
@@ -1131,7 +1137,9 @@ class Person(Base):
 
             class fake_record_sql_queries:  # noqa: N801
                 def __init__(self, *args, **kwargs):
-                    with record_sql_queries(*args, **kwargs) as span:
+                    with record_sql_queries_supporting_streaming(
+                        *args, **kwargs
+                    ) as span:
                         self.span = span
 
                     self.span.start_timestamp = datetime(2024, 1, 1, microsecond=0)
@@ -1144,7 +1152,7 @@ def __exit__(self, type, value, traceback):
                     pass
 
             with mock.patch(
-                "sentry_sdk.integrations.sqlalchemy.record_sql_queries",
+                "sentry_sdk.integrations.sqlalchemy.record_sql_queries_supporting_streaming",
                 fake_record_sql_queries,
             ):
                 assert session.query(Person).first() == bob