Skip to content

Commit 83871a0

Browse files
sentrivanasentrivana
authored
Fix clickhouse breadcrumbs (#3687)
1 parent cae5a32 commit 83871a0

File tree

2 files changed

+53
-56
lines changed

2 files changed

+53
-56
lines changed

‎sentry_sdk/integrations/clickhouse_driver.py

Lines changed: 42 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
ensure_integration_enabled,
1010
)
1111

12-
from typing import TYPE_CHECKING, TypeVar
12+
from typing import TYPE_CHECKING, Any, Dict, TypeVar
1313

1414
# Hack to get new Python features working in older versions
1515
# without introducing a hard dependency on `typing_extensions`
@@ -94,17 +94,17 @@ def _inner(*args: P.args, **kwargs: P.kwargs) -> T:
9494

9595
connection._sentry_span = span # type: ignore[attr-defined]
9696

97-
_set_db_data(span, connection)
98-
99-
if should_send_default_pii():
100-
span.set_attribute("db.query.text", query)
97+
data = _get_db_data(connection)
98+
data["db.query.text"] = query
10199

102100
if query_id:
103-
span.set_attribute("db.query_id", query_id)
101+
data["db.query_id"] = query_id
104102

105103
if params and should_send_default_pii():
106-
connection._sentry_db_params = params
107-
span.set_attribute("db.params", _serialize_span_attribute(params))
104+
data["db.params"] = params
105+
106+
connection._sentry_db_data = data # type: ignore[attr-defined]
107+
_set_on_span(span, data)
108108

109109
# run the original code
110110
ret = f(*args, **kwargs)
@@ -117,69 +117,68 @@ def _inner(*args: P.args, **kwargs: P.kwargs) -> T:
117117
def _wrap_end(f: Callable[P, T]) -> Callable[P, T]:
118118
def _inner_end(*args: P.args, **kwargs: P.kwargs) -> T:
119119
res = f(*args, **kwargs)
120-
instance = args[0]
121-
span = getattr(instance.connection, "_sentry_span", None) # type: ignore[attr-defined]
120+
connection = args[0].connection
122121

122+
span = getattr(connection, "_sentry_span", None) # type: ignore[attr-defined]
123123
if span is not None:
124+
data = getattr(connection, "_sentry_db_data", {})
125+
124126
if res is not None and should_send_default_pii():
127+
data["db.result"] = res
125128
span.set_attribute("db.result", _serialize_span_attribute(res))
126129

127130
with capture_internal_exceptions():
128-
query = span.get_attribute("db.query.text")
131+
query = data.pop("db.query.text", None)
129132
if query:
130-
data = {}
131-
for attr in (
132-
"db.params",
133-
"db.result",
134-
SPANDATA.DB_SYSTEM,
135-
SPANDATA.DB_USER,
136-
SPANDATA.SERVER_ADDRESS,
137-
SPANDATA.SERVER_PORT,
138-
):
139-
if span.get_attribute(attr):
140-
data[attr] = span.get_attribute(attr)
141-
142133
sentry_sdk.add_breadcrumb(
143134
message=query, category="query", data=data
144135
)
145136

146137
span.finish()
147138

139+
try:
140+
del connection._sentry_db_data
141+
del connection._sentry_span
142+
except AttributeError:
143+
pass
144+
148145
return res
149146

150147
return _inner_end
151148

152149

153150
def _wrap_send_data(f: Callable[P, T]) -> Callable[P, T]:
154151
def _inner_send_data(*args: P.args, **kwargs: P.kwargs) -> T:
155-
instance = args[0] # type: clickhouse_driver.client.Client
156-
data = args[2]
157-
span = getattr(instance.connection, "_sentry_span", None)
152+
connection = args[0].connection
153+
db_params_data = args[2]
154+
span = getattr(connection, "_sentry_span", None)
158155

159156
if span is not None:
160-
_set_db_data(span, instance.connection)
157+
data = _get_db_data(connection)
158+
_set_on_span(span, data)
161159

162160
if should_send_default_pii():
163-
db_params = (
164-
getattr(instance.connection, "_sentry_db_params", None) or []
165-
)
166-
db_params.extend(data)
161+
saved_db_data = getattr(connection, "_sentry_db_data", {})
162+
db_params = saved_db_data.get("db.params") or []
163+
db_params.extend(db_params_data)
164+
saved_db_data["db.params"] = db_params
167165
span.set_attribute("db.params", _serialize_span_attribute(db_params))
168-
try:
169-
del instance.connection._sentry_db_params
170-
except AttributeError:
171-
pass
172166

173167
return f(*args, **kwargs)
174168

175169
return _inner_send_data
176170

177171

178-
def _set_db_data(
179-
span: Span, connection: clickhouse_driver.connection.Connection
180-
) -> None:
181-
span.set_attribute(SPANDATA.DB_SYSTEM, "clickhouse")
182-
span.set_attribute(SPANDATA.SERVER_ADDRESS, connection.host)
183-
span.set_attribute(SPANDATA.SERVER_PORT, connection.port)
184-
span.set_attribute(SPANDATA.DB_NAME, connection.database)
185-
span.set_attribute(SPANDATA.DB_USER, connection.user)
172+
def _get_db_data(connection: clickhouse_driver.connection.Connection) -> Dict[str, str]:
173+
return {
174+
SPANDATA.DB_SYSTEM: "clickhouse",
175+
SPANDATA.SERVER_ADDRESS: connection.host,
176+
SPANDATA.SERVER_PORT: connection.port,
177+
SPANDATA.DB_NAME: connection.database,
178+
SPANDATA.DB_USER: connection.user,
179+
}
180+
181+
182+
def _set_on_span(span: Span, data: Dict[str, Any]):
183+
for key, value in data.items():
184+
span.set_attribute(key, _serialize_span_attribute(value))

‎tests/integrations/clickhouse_driver/test_clickhouse_driver.py

Lines changed: 11 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -168,7 +168,7 @@ def test_clickhouse_client_breadcrumbs_with_pii(sentry_init, capture_events) ->
168168
"db.user": "default",
169169
"server.address": "localhost",
170170
"server.port": 9000,
171-
"db.params": '[{"x": 100}]',
171+
"db.params": [{"x": 100}],
172172
},
173173
"message": "INSERT INTO test (x) VALUES",
174174
"type": "default",
@@ -181,7 +181,7 @@ def test_clickhouse_client_breadcrumbs_with_pii(sentry_init, capture_events) ->
181181
"db.user": "default",
182182
"server.address": "localhost",
183183
"server.port": 9000,
184-
"db.params": "[[170], [200]]",
184+
"db.params": [[170], [200]],
185185
},
186186
"message": "INSERT INTO test (x) VALUES",
187187
"type": "default",
@@ -194,8 +194,8 @@ def test_clickhouse_client_breadcrumbs_with_pii(sentry_init, capture_events) ->
194194
"db.user": "default",
195195
"server.address": "localhost",
196196
"server.port": 9000,
197-
"db.result": "[[370]]",
198-
"db.params": '{"minv": 150}',
197+
"db.result": [[370]],
198+
"db.params": {"minv": 150},
199199
},
200200
"message": "SELECT sum(x) FROM test WHERE x > 150",
201201
"type": "default",
@@ -348,9 +348,7 @@ def test_clickhouse_client_spans(
348348
assert event["spans"] == expected_spans
349349

350350

351-
def test_clickhouse_client_spans_with_pii(
352-
sentry_init, capture_events, capture_envelopes
353-
) -> None:
351+
def test_clickhouse_client_spans_with_pii(sentry_init, capture_events) -> None:
354352
sentry_init(
355353
integrations=[ClickhouseDriverIntegration()],
356354
_experiments={"record_sql_params": True},
@@ -620,7 +618,7 @@ def test_clickhouse_dbapi_breadcrumbs_with_pii(sentry_init, capture_events) -> N
620618
"db.user": "default",
621619
"server.address": "localhost",
622620
"server.port": 9000,
623-
"db.result": "[[], []]",
621+
"db.result": [[], []],
624622
},
625623
"message": "DROP TABLE IF EXISTS test",
626624
"type": "default",
@@ -633,7 +631,7 @@ def test_clickhouse_dbapi_breadcrumbs_with_pii(sentry_init, capture_events) -> N
633631
"db.user": "default",
634632
"server.address": "localhost",
635633
"server.port": 9000,
636-
"db.result": "[[], []]",
634+
"db.result": [[], []],
637635
},
638636
"message": "CREATE TABLE test (x Int32) ENGINE = Memory",
639637
"type": "default",
@@ -646,7 +644,7 @@ def test_clickhouse_dbapi_breadcrumbs_with_pii(sentry_init, capture_events) -> N
646644
"db.user": "default",
647645
"server.address": "localhost",
648646
"server.port": 9000,
649-
"db.params": '[{"x": 100}]',
647+
"db.params": [{"x": 100}],
650648
},
651649
"message": "INSERT INTO test (x) VALUES",
652650
"type": "default",
@@ -659,7 +657,7 @@ def test_clickhouse_dbapi_breadcrumbs_with_pii(sentry_init, capture_events) -> N
659657
"db.user": "default",
660658
"server.address": "localhost",
661659
"server.port": 9000,
662-
"db.params": "[[170], [200]]",
660+
"db.params": [[170], [200]],
663661
},
664662
"message": "INSERT INTO test (x) VALUES",
665663
"type": "default",
@@ -672,8 +670,8 @@ def test_clickhouse_dbapi_breadcrumbs_with_pii(sentry_init, capture_events) -> N
672670
"db.user": "default",
673671
"server.address": "localhost",
674672
"server.port": 9000,
675-
"db.params": '{"minv": 150}',
676-
"db.result": '[[["370"]], [["\'sum(x)\'", "\'Int64\'"]]]',
673+
"db.params": {"minv": 150},
674+
"db.result": [[["370"]], [["'sum(x)'", "'Int64'"]]],
677675
},
678676
"message": "SELECT sum(x) FROM test WHERE x > 150",
679677
"type": "default",

0 commit comments

Comments
 (0)