feat(clickhouse_driver): Support span streaming

sentry_sdk/integrations/clickhouse_driver.py

@@ -4,15 +4,17 @@
 from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.integrations import DidNotEnable, Integration, _check_minimum_version
 from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.traces import StreamedSpan
 from sentry_sdk.tracing import Span
-from sentry_sdk.utils import capture_internal_exceptions, ensure_integration_enabled
+from sentry_sdk.tracing_utils import has_span_streaming_enabled
+from sentry_sdk.utils import capture_internal_exceptions
 
 # Hack to get new Python features working in older versions
 # without introducing a hard dependency on `typing_extensions`
 # from: https://stackoverflow.com/a/71944042/300572
 if TYPE_CHECKING:
     from collections.abc import Iterator
-    from typing import Any, Callable, ParamSpec
+    from typing import Any, Callable, ParamSpec, Union
 else:
     # Fake ParamSpec
     class ParamSpec:
@@ -70,30 +72,42 @@ def setup_once() -> None:
 
 
 def _wrap_start(f: "Callable[P, T]") -> "Callable[P, T]":
-    @ensure_integration_enabled(ClickhouseDriverIntegration, f)
     def _inner(*args: "P.args", **kwargs: "P.kwargs") -> "T":
+        client = sentry_sdk.get_client()
+        if client.get_integration(ClickhouseDriverIntegration) is None:
+            return f(*args, **kwargs)
+
         connection = args[0]
         query = args[1]
         query_id = args[2] if len(args) > 2 else kwargs.get("query_id")
         params = args[3] if len(args) > 3 else kwargs.get("params")
 
-        span = sentry_sdk.start_span(
-            op=OP.DB,
-            name=query,
-            origin=ClickhouseDriverIntegration.origin,
-        )
+        if has_span_streaming_enabled(client.options):
+            span = sentry_sdk.traces.start_span(
+                name=query,  # type: ignore
+                attributes={
+                    "sentry.op": OP.DB,
+                    "sentry.origin": ClickhouseDriverIntegration.origin,
+                },
+            )
+        else:
+            span = sentry_sdk.start_span(
+                op=OP.DB,
+                name=query,
+                origin=ClickhouseDriverIntegration.origin,
+            )
 
-        connection._sentry_span = span  # type: ignore[attr-defined]
+            span.set_data("query", query)
 
-        _set_db_data(span, connection)
+            if query_id:
+                span.set_data("db.query_id", query_id)
 
-        span.set_data("query", query)
+            if params and should_send_default_pii():
+                span.set_data("db.params", params)
 
-        if query_id:
-            span.set_data("db.query_id", query_id)
+        connection._sentry_span = span  # type: ignore[attr-defined]
 
-        if params and should_send_default_pii():
-            span.set_data("db.params", params)
+        _set_db_data(span, connection)
 
         # run the original code
         ret = f(*args, **kwargs)
@@ -109,7 +123,12 @@ def _inner_end(*args: "P.args", **kwargs: "P.kwargs") -> "T":
         instance = args[0]
         span = getattr(instance.connection, "_sentry_span", None)  # type: ignore[attr-defined]
 
-        if span is not None:
+        if span is None:
+            return res
+
+        if isinstance(span, StreamedSpan):
+            span.end()
+        else:
             if res is not None and should_send_default_pii():
                 span.set_data("db.result", res)
 
@@ -133,6 +152,12 @@ def _inner_send_data(  # type: ignore[no-untyped-def] # clickhouse-driver does n
     ):
         span = getattr(self.connection, "_sentry_span", None)
 
+        if isinstance(span, StreamedSpan):
+            _set_db_data(span, self.connection)
+            return original_send_data(
+                self, sample_block, data, types_check, columnar, *args, **kwargs
+            )
+
         if span is not None:
             _set_db_data(span, self.connection)
 
@@ -165,10 +190,19 @@ def wrapped_generator() -> "Iterator[Any]":
     Client.send_data = _inner_send_data
 
 
-def _set_db_data(span: "Span", connection: "Connection") -> None:
-    span.set_data(SPANDATA.DB_SYSTEM, "clickhouse")
-    span.set_data(SPANDATA.DB_DRIVER_NAME, "clickhouse-driver")
-    span.set_data(SPANDATA.SERVER_ADDRESS, connection.host)
-    span.set_data(SPANDATA.SERVER_PORT, connection.port)
-    span.set_data(SPANDATA.DB_NAME, connection.database)
-    span.set_data(SPANDATA.DB_USER, connection.user)
+def _set_db_data(span: "Union[Span, StreamedSpan]", connection: "Connection") -> None:
+    if isinstance(span, StreamedSpan):
+        span.set_attribute(SPANDATA.DB_SYSTEM_NAME, "clickhouse")
+        span.set_attribute(SPANDATA.DB_NAMESPACE, connection.database)
+
+        set_on_span = span.set_attribute
+    else:
+        span.set_data(SPANDATA.DB_SYSTEM, "clickhouse")
+        span.set_data(SPANDATA.DB_NAME, connection.database)
+
+        set_on_span = span.set_data
+
+    set_on_span(SPANDATA.DB_DRIVER_NAME, "clickhouse-driver")
+    set_on_span(SPANDATA.SERVER_ADDRESS, connection.host)
+    set_on_span(SPANDATA.SERVER_PORT, connection.port)
+    set_on_span(SPANDATA.DB_USER, connection.user)