Skip to content

Commit 9d2b826

Browse files
feat(clickhouse_driver): Set breadcrumbs in the streaming trace lifecycle (#7325)
Set the attributes used by streamed spans as breadcrumbs when the streaming trace lifecycle is enabled. Adds tests with the `_span_streaming` suffix based on existing tests.
1 parent 4ae8604 commit 9d2b826

2 files changed

Lines changed: 968 additions & 58 deletions

File tree

sentry_sdk/integrations/clickhouse_driver.py

Lines changed: 38 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
# from: https://stackoverflow.com/a/71944042/300572
1616
if TYPE_CHECKING:
1717
from collections.abc import Iterator
18-
from typing import Any, Callable, ParamSpec, Union
18+
from typing import Any, Callable, Optional, ParamSpec, Union
1919
else:
2020
# Fake ParamSpec
2121
class ParamSpec:
@@ -79,7 +79,7 @@ def _inner(*args: "P.args", **kwargs: "P.kwargs") -> "T":
7979
if client.get_integration(ClickhouseDriverIntegration) is None:
8080
return f(*args, **kwargs)
8181

82-
connection = args[0]
82+
connection: "Connection" = args[0]
8383
query = args[1]
8484
query_id = args[2] if len(args) > 2 else kwargs.get("query_id")
8585
params = args[3] if len(args) > 3 else kwargs.get("params")
@@ -95,6 +95,16 @@ def _inner(*args: "P.args", **kwargs: "P.kwargs") -> "T":
9595
SPANDATA.DB_QUERY_TEXT: str(query),
9696
},
9797
)
98+
99+
connection._query = query
100+
connection._breadcrumb_data = {
101+
SPANDATA.DB_SYSTEM: "clickhouse",
102+
SPANDATA.DB_NAME: connection.database,
103+
SPANDATA.DB_DRIVER_NAME: "clickhouse-driver",
104+
SPANDATA.SERVER_ADDRESS: connection.host,
105+
SPANDATA.SERVER_PORT: connection.port,
106+
SPANDATA.DB_USER: connection.user,
107+
}
98108
else:
99109
span = sentry_sdk.start_span(
100110
op=OP.DB,
@@ -114,7 +124,7 @@ def _inner(*args: "P.args", **kwargs: "P.kwargs") -> "T":
114124
elif should_send_default_pii():
115125
span.set_data("db.params", params)
116126

117-
connection._sentry_span = span # type: ignore[attr-defined]
127+
connection._sentry_span = span
118128

119129
if span is not None:
120130
_set_db_data(span, connection)
@@ -130,8 +140,31 @@ def _inner(*args: "P.args", **kwargs: "P.kwargs") -> "T":
130140
def _wrap_end(f: "Callable[P, T]") -> "Callable[P, T]":
131141
def _inner_end(*args: "P.args", **kwargs: "P.kwargs") -> "T":
132142
res = f(*args, **kwargs)
133-
instance = args[0]
134-
span = getattr(instance.connection, "_sentry_span", None) # type: ignore[attr-defined]
143+
instance: "Client" = args[0]
144+
145+
query = getattr(instance.connection, "_query", None)
146+
breadcrumb_data: "Optional[dict[str, Any]]" = getattr(
147+
instance.connection, "_breadcrumb_data", None
148+
)
149+
150+
if query is not None and breadcrumb_data is not None:
151+
client_options = sentry_sdk.get_client().options
152+
if (
153+
has_data_collection_enabled(client_options)
154+
and client_options["data_collection"]["database_query_data"]
155+
) or (
156+
not has_data_collection_enabled(client_options)
157+
and should_send_default_pii()
158+
):
159+
breadcrumb_data = {"db.result": res, **breadcrumb_data}
160+
161+
sentry_sdk.get_isolation_scope().add_breadcrumb(
162+
message=query,
163+
category="query",
164+
data={"db.result": res, **breadcrumb_data},
165+
)
166+
167+
span = getattr(instance.connection, "_sentry_span", None)
135168

136169
if span is None:
137170
return res

0 commit comments

Comments
 (0)