Fix clickhouse-driver integration spans (#3638)

Author: sentrivana

MIGRATION_GUIDE.md

@@ -18,6 +18,7 @@ Looking to upgrade from Sentry SDK 2.x to 3.x? Here's a comprehensive list of wh
 - A `Profile` object does not have a `.hub` property anymore.
 - `sentry_sdk.continue_trace` no longer returns a `Transaction` and is now a context manager.
 - Redis integration: In Redis pipeline spans there is no `span["data"]["redis.commands"]` that contains a dict `{"count": 3, "first_ten": ["cmd1", "cmd2", ...]}` but instead `span["data"]["redis.commands.count"]` (containing `3`) and `span["data"]["redis.commands.first_ten"]` (containing `["cmd1", "cmd2", ...]`).
+- clickhouse-driver integration: The query is now available under the `db.query.text` span attribute (only if `send_default_pii` is `True`).
 
 ### Removed
 

sentry_sdk/integrations/clickhouse_driver.py

@@ -91,13 +91,15 @@ def _inner(*args: P.args, **kwargs: P.kwargs) -> T:
 
         _set_db_data(span, connection)
 
-        span.set_data("query", query)
+        if should_send_default_pii():
+            span.set_attribute("db.query.text", query)
 
         if query_id:
-            span.set_data("db.query_id", query_id)
+            span.set_attribute("db.query_id", query_id)
 
         if params and should_send_default_pii():
-            span.set_data("db.params", params)
+            connection._sentry_db_params = params
+            span.set_attribute("db.params", str(params))
 
         # run the original code
         ret = f(*args, **kwargs)
@@ -115,12 +117,26 @@ def _inner_end(*args: P.args, **kwargs: P.kwargs) -> T:
 
         if span is not None:
             if res is not None and should_send_default_pii():
-                span.set_data("db.result", res)
+                span.set_attribute("db.result", str(res))
 
             with capture_internal_exceptions():
-                span.scope.add_breadcrumb(
-                    message=span._data.pop("query"), category="query", data=span._data
-                )
+                query = span.get_attribute("db.query.text")
+                if query:
+                    data = {}
+                    for attr in (
+                        "db.params",
+                        "db.result",
+                        SPANDATA.DB_SYSTEM,
+                        SPANDATA.DB_USER,
+                        SPANDATA.SERVER_ADDRESS,
+                        SPANDATA.SERVER_PORT,
+                    ):
+                        if span.get_attribute(attr):
+                            data[attr] = span.get_attribute(attr)
+
+                    sentry_sdk.add_breadcrumb(
+                        message=query, category="query", data=data
+                    )
 
             span.finish()
 
@@ -139,9 +155,15 @@ def _inner_send_data(*args: P.args, **kwargs: P.kwargs) -> T:
             _set_db_data(span, instance.connection)
 
             if should_send_default_pii():
-                db_params = span._data.get("db.params", [])
+                db_params = (
+                    getattr(instance.connection, "_sentry_db_params", None) or []
+                )
                 db_params.extend(data)
-                span.set_data("db.params", db_params)
+                span.set_attribute("db.params", str(db_params))
+                try:
+                    del instance.connection._sentry_db_params
+                except AttributeError:
+                    pass
 
         return f(*args, **kwargs)
 
@@ -151,8 +173,8 @@ def _inner_send_data(*args: P.args, **kwargs: P.kwargs) -> T:
 def _set_db_data(
     span: Span, connection: clickhouse_driver.connection.Connection
 ) -> None:
-    span.set_data(SPANDATA.DB_SYSTEM, "clickhouse")
-    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)
+    span.set_attribute(SPANDATA.DB_SYSTEM, "clickhouse")
+    span.set_attribute(SPANDATA.SERVER_ADDRESS, connection.host)
+    span.set_attribute(SPANDATA.SERVER_PORT, connection.port)
+    span.set_attribute(SPANDATA.DB_NAME, connection.database)
+    span.set_attribute(SPANDATA.DB_USER, connection.user)

sentry_sdk/tracing.py

@@ -1279,18 +1279,12 @@ def __exit__(self, ty, value, tb):
         self.finish()
         context.detach(self._ctx_token)
 
-    def _get_attribute(self, name):
-        # type: (str) -> Optional[Any]
-        if not isinstance(self._otel_span, ReadableSpan):
-            return None
-        return self._otel_span.attributes.get(name)
-
     @property
     def description(self):
         # type: () -> Optional[str]
         from sentry_sdk.integrations.opentelemetry.consts import SentrySpanAttribute
 
-        return self._get_attribute(SentrySpanAttribute.DESCRIPTION)
+        return self.get_attribute(SentrySpanAttribute.DESCRIPTION)
 
     @description.setter
     def description(self, value):
@@ -1305,7 +1299,7 @@ def origin(self):
         # type: () -> Optional[str]
         from sentry_sdk.integrations.opentelemetry.consts import SentrySpanAttribute
 
-        return self._get_attribute(SentrySpanAttribute.ORIGIN)
+        return self.get_attribute(SentrySpanAttribute.ORIGIN)
 
     @origin.setter
     def origin(self, value):
@@ -1378,7 +1372,7 @@ def op(self):
         # type: () -> Optional[str]
         from sentry_sdk.integrations.opentelemetry.consts import SentrySpanAttribute
 
-        return self._get_attribute(SentrySpanAttribute.OP)
+        return self.get_attribute(SentrySpanAttribute.OP)
 
     @op.setter
     def op(self, value):
@@ -1393,7 +1387,7 @@ def name(self):
         # type: () -> Optional[str]
         from sentry_sdk.integrations.opentelemetry.consts import SentrySpanAttribute
 
-        return self._get_attribute(SentrySpanAttribute.NAME)
+        return self.get_attribute(SentrySpanAttribute.NAME)
 
     @name.setter
     def name(self, value):
@@ -1509,6 +1503,12 @@ def set_data(self, key, value):
         # TODO-neel-potel we cannot add dicts here
         self.set_attribute(key, value)
 
+    def get_attribute(self, name):
+        # type: (str) -> Optional[Any]
+        if not isinstance(self._otel_span, ReadableSpan):
+            return None
+        return self._otel_span.attributes.get(name)
+
     def set_attribute(self, key, value):
         # type: (str, Any) -> None
         self._otel_span.set_attribute(key, value)