diff --git a/sentry_sdk/integrations/clickhouse_driver.py b/sentry_sdk/integrations/clickhouse_driver.py index 8596ed170f..1f0ee15d9f 100644 --- a/sentry_sdk/integrations/clickhouse_driver.py +++ b/sentry_sdk/integrations/clickhouse_driver.py @@ -146,40 +146,47 @@ def _inner_send_data( **kwargs: Any, ) -> Any: span = getattr(self.connection, "_sentry_span", None) + if span is None: + return original_send_data( + self, sample_block, data, types_check, columnar, *args, **kwargs + ) - if span is not None: - data = _get_db_data(self.connection) - _set_on_span(span, data) - - if should_send_default_pii(): - saved_db_data: dict[str, Any] = getattr( - self.connection, "_sentry_db_data", {} - ) - db_params: list[Any] = saved_db_data.get("db.params") or [] + db_data = _get_db_data(self.connection) + _set_on_span(span, db_data) - if isinstance(data, (list, tuple)): - db_params.extend(data) + saved_db_data: dict[str, Any] = getattr(self.connection, "_sentry_db_data", {}) + db_params: list[Any] = saved_db_data.get("db.params") or [] - else: # data is a generic iterator - orig_data = data + if should_send_default_pii(): + if isinstance(data, (list, tuple)): + db_params.extend(data) - # Wrap the generator to add items to db.params as they are yielded. - # This allows us to send the params to Sentry without needing to allocate - # memory for the entire generator at once. - def wrapped_generator() -> "Iterator[Any]": - for item in orig_data: - db_params.append(item) - yield item + else: # data is a generic iterator + orig_data = data - # Replace the original iterator with the wrapped one. - data = wrapped_generator() + # Wrap the generator to add items to db.params as they are yielded. + # This allows us to send the params to Sentry without needing to allocate + # memory for the entire generator at once. + def wrapped_generator() -> "Iterator[Any]": + for item in orig_data: + db_params.append(item) + yield item - span.set_attribute("db.params", _serialize_span_attribute(db_params)) + # Replace the original iterator with the wrapped one. + data = wrapped_generator() - return original_send_data( + rv = original_send_data( self, sample_block, data, types_check, columnar, *args, **kwargs ) + if should_send_default_pii() and db_params: + # need to do this after the original function call to make sure + # db_params is populated correctly + saved_db_data["db.params"] = db_params + span.set_attribute("db.params", _serialize_span_attribute(db_params)) + + return rv + clickhouse_driver.client.Client.send_data = _inner_send_data diff --git a/tests/integrations/clickhouse_driver/test_clickhouse_driver.py b/tests/integrations/clickhouse_driver/test_clickhouse_driver.py index a9ec386c22..3ad6ce93fd 100644 --- a/tests/integrations/clickhouse_driver/test_clickhouse_driver.py +++ b/tests/integrations/clickhouse_driver/test_clickhouse_driver.py @@ -383,7 +383,7 @@ def test_clickhouse_spans_with_generator(sentry_init, capture_events): span for span in spans if span["description"] == "INSERT INTO test (x) VALUES" ] - assert span["data"]["db.params"] == [{"x": 0}, {"x": 1}, {"x": 2}] + assert span["data"]["db.params"] == '[{"x": 0}, {"x": 1}, {"x": 2}]' def test_clickhouse_client_spans_with_pii(