Fix clickhouse logic for generator case

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
 
 

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(