Skip to content

Commit 3ef02a1

Browse files
fix(clickhouse): Don't eat the generator data (#4669)
Currently, the Clickhouse integration consumes any data passed as a generator when reading it for insertion as `db_params`. Instead, since generators cannot be cloned, we need to wrap the generator to add the params as we iterate over it. Fixes #4657 <!-- Describe your PR here --> --- Thank you for contributing to `sentry-python`! Please add tests to validate your changes, and lint your code using `tox -e linters`. Running the test suite on your PR might require maintainer approval.
1 parent b73f876 commit 3ef02a1

File tree

2 files changed

+52
-2
lines changed

2 files changed

+52
-2
lines changed

sentry_sdk/integrations/clickhouse_driver.py

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,8 @@
1111
# without introducing a hard dependency on `typing_extensions`
1212
# from: https://stackoverflow.com/a/71944042/300572
1313
if TYPE_CHECKING:
14-
from typing import ParamSpec, Callable
14+
from collections.abc import Iterator
15+
from typing import Any, ParamSpec, Callable
1516
else:
1617
# Fake ParamSpec
1718
class ParamSpec:
@@ -139,7 +140,24 @@ def _inner_send_data( # type: ignore[no-untyped-def] # clickhouse-driver does n
139140

140141
if should_send_default_pii():
141142
db_params = span._data.get("db.params", [])
142-
db_params.extend(data)
143+
144+
if isinstance(data, (list, tuple)):
145+
db_params.extend(data)
146+
147+
else: # data is a generic iterator
148+
orig_data = data
149+
150+
# Wrap the generator to add items to db.params as they are yielded.
151+
# This allows us to send the params to Sentry without needing to allocate
152+
# memory for the entire generator at once.
153+
def wrapped_generator() -> "Iterator[Any]":
154+
for item in orig_data:
155+
db_params.append(item)
156+
yield item
157+
158+
# Replace the original iterator with the wrapped one.
159+
data = wrapped_generator()
160+
143161
span.set_data("db.params", db_params)
144162

145163
return original_send_data(

tests/integrations/clickhouse_driver/test_clickhouse_driver.py

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -342,6 +342,38 @@ def test_clickhouse_client_spans(
342342
assert event["spans"] == expected_spans
343343

344344

345+
def test_clickhouse_spans_with_generator(sentry_init, capture_events):
346+
sentry_init(
347+
integrations=[ClickhouseDriverIntegration()],
348+
send_default_pii=True,
349+
traces_sample_rate=1.0,
350+
)
351+
events = capture_events()
352+
353+
# Use a generator to test that the integration obtains values from the generator,
354+
# without consuming the generator.
355+
values = ({"x": i} for i in range(3))
356+
357+
with start_transaction(name="test_clickhouse_transaction"):
358+
client = Client("localhost")
359+
client.execute("DROP TABLE IF EXISTS test")
360+
client.execute("CREATE TABLE test (x Int32) ENGINE = Memory")
361+
client.execute("INSERT INTO test (x) VALUES", values)
362+
res = client.execute("SELECT x FROM test")
363+
364+
# Verify that the integration did not consume the generator
365+
assert res == [(0,), (1,), (2,)]
366+
367+
(event,) = events
368+
spans = event["spans"]
369+
370+
[span] = [
371+
span for span in spans if span["description"] == "INSERT INTO test (x) VALUES"
372+
]
373+
374+
assert span["data"]["db.params"] == [{"x": 0}, {"x": 1}, {"x": 2}]
375+
376+
345377
def test_clickhouse_client_spans_with_pii(
346378
sentry_init, capture_events, capture_envelopes
347379
) -> None:

0 commit comments

Comments
 (0)