diff --git a/sentry_sdk/integrations/clickhouse_driver.py b/sentry_sdk/integrations/clickhouse_driver.py index 9583b25ca9..4e9901acaf 100644 --- a/sentry_sdk/integrations/clickhouse_driver.py +++ b/sentry_sdk/integrations/clickhouse_driver.py @@ -8,7 +8,7 @@ from sentry_sdk.traces import StreamedSpan from sentry_sdk.tracing import Span from sentry_sdk.tracing_utils import has_span_streaming_enabled -from sentry_sdk.utils import capture_internal_exceptions +from sentry_sdk.utils import capture_internal_exceptions, has_data_collection_enabled # Hack to get new Python features working in older versions # without introducing a hard dependency on `typing_extensions` @@ -107,8 +107,12 @@ def _inner(*args: "P.args", **kwargs: "P.kwargs") -> "T": if query_id: span.set_data("db.query_id", query_id) - if params and should_send_default_pii(): - span.set_data("db.params", params) + if params: + if has_data_collection_enabled(client.options): + if client.options["data_collection"]["database_query_data"]: + span.set_data("db.params", params) + elif should_send_default_pii(): + span.set_data("db.params", params) connection._sentry_span = span # type: ignore[attr-defined] @@ -135,8 +139,13 @@ def _inner_end(*args: "P.args", **kwargs: "P.kwargs") -> "T": if isinstance(span, StreamedSpan): span.end() else: - if res is not None and should_send_default_pii(): - span.set_data("db.result", res) + if res is not None: + client_options = sentry_sdk.get_client().options + if has_data_collection_enabled(client_options): + if client_options["data_collection"]["database_query_data"]: + span.set_data("db.result", res) + elif should_send_default_pii(): + span.set_data("db.result", res) with capture_internal_exceptions(): span.scope.add_breadcrumb( @@ -167,7 +176,29 @@ def _inner_send_data( # type: ignore[no-untyped-def] # clickhouse-driver does n if span is not None: _set_db_data(span, self.connection) - if should_send_default_pii(): + client_options = sentry_sdk.get_client().options + if has_data_collection_enabled(client_options): + if client_options["data_collection"]["database_query_data"]: + db_params = span._data.get("db.params", []) + if isinstance(data, (list, tuple)): + db_params.extend(data) + + else: # data is a generic iterator + orig_data = 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 + + # Replace the original iterator with the wrapped one. + data = wrapped_generator() + + span.set_data("db.params", db_params) + elif should_send_default_pii(): db_params = span._data.get("db.params", []) if isinstance(data, (list, tuple)): diff --git a/tests/integrations/clickhouse_driver/test_clickhouse_driver.py b/tests/integrations/clickhouse_driver/test_clickhouse_driver.py index 7e735ec160..d19f630c4b 100644 --- a/tests/integrations/clickhouse_driver/test_clickhouse_driver.py +++ b/tests/integrations/clickhouse_driver/test_clickhouse_driver.py @@ -229,6 +229,510 @@ def test_clickhouse_client_breadcrumbs_with_pii(sentry_init, capture_events) -> assert event["breadcrumbs"]["values"] == expected_breadcrumbs +def test_clickhouse_client_breadcrumbs_with_data_collection( + sentry_init, capture_events +) -> None: + sentry_init( + integrations=[ClickhouseDriverIntegration()], + _experiments={"data_collection": {"database_query_data": True}}, + ) + events = capture_events() + + client = Client("localhost") + client.execute("DROP TABLE IF EXISTS test") + client.execute("CREATE TABLE test (x Int32) ENGINE = Memory") + client.execute("INSERT INTO test (x) VALUES", [{"x": 100}]) + client.execute("INSERT INTO test (x) VALUES", [[170], [200]]) + + res = client.execute("SELECT sum(x) FROM test WHERE x > %(minv)i", {"minv": 150}) + assert res[0][0] == 370 + + capture_message("hi") + + (event,) = events + + expected_breadcrumbs = [ + { + "category": "query", + "data": { + "db.system": "clickhouse", + "db.name": "", + "db.user": "default", + "server.address": "localhost", + "server.port": 9000, + "db.result": [], + }, + "message": "DROP TABLE IF EXISTS test", + "type": "default", + }, + { + "category": "query", + "data": { + "db.system": "clickhouse", + "db.name": "", + "db.user": "default", + "server.address": "localhost", + "server.port": 9000, + "db.result": [], + }, + "message": "CREATE TABLE test (x Int32) ENGINE = Memory", + "type": "default", + }, + { + "category": "query", + "data": { + "db.system": "clickhouse", + "db.name": "", + "db.user": "default", + "server.address": "localhost", + "server.port": 9000, + "db.params": [{"x": 100}], + }, + "message": "INSERT INTO test (x) VALUES", + "type": "default", + }, + { + "category": "query", + "data": { + "db.system": "clickhouse", + "db.name": "", + "db.user": "default", + "server.address": "localhost", + "server.port": 9000, + "db.params": [[170], [200]], + }, + "message": "INSERT INTO test (x) VALUES", + "type": "default", + }, + { + "category": "query", + "data": { + "db.system": "clickhouse", + "db.name": "", + "db.user": "default", + "server.address": "localhost", + "server.port": 9000, + "db.result": [[370]], + "db.params": {"minv": 150}, + }, + "message": "SELECT sum(x) FROM test WHERE x > 150", + "type": "default", + }, + ] + + if not EXPECT_PARAMS_IN_SELECT: + expected_breadcrumbs[-1]["data"].pop("db.params", None) + + for crumb in expected_breadcrumbs: + crumb["data"] = ApproxDict(crumb["data"]) + + for crumb in event["breadcrumbs"]["values"]: + crumb.pop("timestamp", None) + + assert event["breadcrumbs"]["values"] == expected_breadcrumbs + + +def test_clickhouse_client_breadcrumbs_with_data_collection_disabled( + sentry_init, capture_events +) -> None: + sentry_init( + integrations=[ClickhouseDriverIntegration()], + _experiments={"data_collection": {"database_query_data": False}}, + ) + events = capture_events() + + client = Client("localhost") + client.execute("DROP TABLE IF EXISTS test") + client.execute("CREATE TABLE test (x Int32) ENGINE = Memory") + client.execute("INSERT INTO test (x) VALUES", [{"x": 100}]) + client.execute("INSERT INTO test (x) VALUES", [[170], [200]]) + + res = client.execute("SELECT sum(x) FROM test WHERE x > %(minv)i", {"minv": 150}) + assert res[0][0] == 370 + + capture_message("hi") + + (event,) = events + + expected_breadcrumbs = [ + { + "category": "query", + "data": { + "db.system": "clickhouse", + "db.name": "", + "db.user": "default", + "server.address": "localhost", + "server.port": 9000, + }, + "message": "DROP TABLE IF EXISTS test", + "type": "default", + }, + { + "category": "query", + "data": { + "db.system": "clickhouse", + "db.name": "", + "db.user": "default", + "server.address": "localhost", + "server.port": 9000, + }, + "message": "CREATE TABLE test (x Int32) ENGINE = Memory", + "type": "default", + }, + { + "category": "query", + "data": { + "db.system": "clickhouse", + "db.name": "", + "db.user": "default", + "server.address": "localhost", + "server.port": 9000, + }, + "message": "INSERT INTO test (x) VALUES", + "type": "default", + }, + { + "category": "query", + "data": { + "db.system": "clickhouse", + "db.name": "", + "db.user": "default", + "server.address": "localhost", + "server.port": 9000, + }, + "message": "INSERT INTO test (x) VALUES", + "type": "default", + }, + { + "category": "query", + "data": { + "db.system": "clickhouse", + "db.name": "", + "db.user": "default", + "server.address": "localhost", + "server.port": 9000, + }, + "message": "SELECT sum(x) FROM test WHERE x > 150", + "type": "default", + }, + ] + + for crumb in expected_breadcrumbs: + crumb["data"] = ApproxDict(crumb["data"]) + + for crumb in event["breadcrumbs"]["values"]: + crumb.pop("timestamp", None) + + assert event["breadcrumbs"]["values"] == expected_breadcrumbs + + # ApproxDict is subset matching, so explicitly assert params/result absence + for crumb in event["breadcrumbs"]["values"]: + assert "db.params" not in crumb["data"] + assert "db.result" not in crumb["data"] + + +def test_clickhouse_client_breadcrumbs_data_collection_overrides_pii( + sentry_init, capture_events +) -> None: + sentry_init( + integrations=[ClickhouseDriverIntegration()], + send_default_pii=True, + _experiments={"data_collection": {"database_query_data": False}}, + ) + events = capture_events() + + client = Client("localhost") + client.execute("DROP TABLE IF EXISTS test") + client.execute("CREATE TABLE test (x Int32) ENGINE = Memory") + client.execute("INSERT INTO test (x) VALUES", [{"x": 100}]) + client.execute("INSERT INTO test (x) VALUES", [[170], [200]]) + + res = client.execute("SELECT sum(x) FROM test WHERE x > %(minv)i", {"minv": 150}) + assert res[0][0] == 370 + + capture_message("hi") + + (event,) = events + + expected_breadcrumbs = [ + { + "category": "query", + "data": { + "db.system": "clickhouse", + "db.name": "", + "db.user": "default", + "server.address": "localhost", + "server.port": 9000, + }, + "message": "DROP TABLE IF EXISTS test", + "type": "default", + }, + { + "category": "query", + "data": { + "db.system": "clickhouse", + "db.name": "", + "db.user": "default", + "server.address": "localhost", + "server.port": 9000, + }, + "message": "CREATE TABLE test (x Int32) ENGINE = Memory", + "type": "default", + }, + { + "category": "query", + "data": { + "db.system": "clickhouse", + "db.name": "", + "db.user": "default", + "server.address": "localhost", + "server.port": 9000, + }, + "message": "INSERT INTO test (x) VALUES", + "type": "default", + }, + { + "category": "query", + "data": { + "db.system": "clickhouse", + "db.name": "", + "db.user": "default", + "server.address": "localhost", + "server.port": 9000, + }, + "message": "INSERT INTO test (x) VALUES", + "type": "default", + }, + { + "category": "query", + "data": { + "db.system": "clickhouse", + "db.name": "", + "db.user": "default", + "server.address": "localhost", + "server.port": 9000, + }, + "message": "SELECT sum(x) FROM test WHERE x > 150", + "type": "default", + }, + ] + + for crumb in expected_breadcrumbs: + crumb["data"] = ApproxDict(crumb["data"]) + + for crumb in event["breadcrumbs"]["values"]: + crumb.pop("timestamp", None) + + assert event["breadcrumbs"]["values"] == expected_breadcrumbs + + # ApproxDict is subset matching, so explicitly assert params/result absence + for crumb in event["breadcrumbs"]["values"]: + assert "db.params" not in crumb["data"] + assert "db.result" not in crumb["data"] + + +def test_clickhouse_client_breadcrumbs_with_data_collection_default( + sentry_init, capture_events +) -> None: + sentry_init( + integrations=[ClickhouseDriverIntegration()], + _experiments={"data_collection": {}}, + ) + events = capture_events() + + client = Client("localhost") + client.execute("DROP TABLE IF EXISTS test") + client.execute("CREATE TABLE test (x Int32) ENGINE = Memory") + client.execute("INSERT INTO test (x) VALUES", [{"x": 100}]) + client.execute("INSERT INTO test (x) VALUES", [[170], [200]]) + + res = client.execute("SELECT sum(x) FROM test WHERE x > %(minv)i", {"minv": 150}) + assert res[0][0] == 370 + + capture_message("hi") + + (event,) = events + + expected_breadcrumbs = [ + { + "category": "query", + "data": { + "db.system": "clickhouse", + "db.name": "", + "db.user": "default", + "server.address": "localhost", + "server.port": 9000, + "db.result": [], + }, + "message": "DROP TABLE IF EXISTS test", + "type": "default", + }, + { + "category": "query", + "data": { + "db.system": "clickhouse", + "db.name": "", + "db.user": "default", + "server.address": "localhost", + "server.port": 9000, + "db.result": [], + }, + "message": "CREATE TABLE test (x Int32) ENGINE = Memory", + "type": "default", + }, + { + "category": "query", + "data": { + "db.system": "clickhouse", + "db.name": "", + "db.user": "default", + "server.address": "localhost", + "server.port": 9000, + "db.params": [{"x": 100}], + }, + "message": "INSERT INTO test (x) VALUES", + "type": "default", + }, + { + "category": "query", + "data": { + "db.system": "clickhouse", + "db.name": "", + "db.user": "default", + "server.address": "localhost", + "server.port": 9000, + "db.params": [[170], [200]], + }, + "message": "INSERT INTO test (x) VALUES", + "type": "default", + }, + { + "category": "query", + "data": { + "db.system": "clickhouse", + "db.name": "", + "db.user": "default", + "server.address": "localhost", + "server.port": 9000, + "db.result": [[370]], + "db.params": {"minv": 150}, + }, + "message": "SELECT sum(x) FROM test WHERE x > 150", + "type": "default", + }, + ] + + if not EXPECT_PARAMS_IN_SELECT: + expected_breadcrumbs[-1]["data"].pop("db.params", None) + + for crumb in expected_breadcrumbs: + crumb["data"] = ApproxDict(crumb["data"]) + + for crumb in event["breadcrumbs"]["values"]: + crumb.pop("timestamp", None) + + assert event["breadcrumbs"]["values"] == expected_breadcrumbs + + +def test_clickhouse_client_span_streaming_with_data_collection( + sentry_init, capture_items +) -> None: + # Streamed spans never carry db.params/db.result, regardless of the + # data_collection configuration + sentry_init( + integrations=[ClickhouseDriverIntegration()], + traces_sample_rate=1.0, + _experiments={ + "trace_lifecycle": "stream", + "data_collection": {"database_query_data": True}, + }, + ) + items = capture_items("span") + + with sentry_sdk.traces.start_span(name="custom parent"): + client = Client("localhost") + client.execute("DROP TABLE IF EXISTS test") + client.execute("CREATE TABLE test (x Int32) ENGINE = Memory") + client.execute("INSERT INTO test (x) VALUES", [{"x": 100}]) + client.execute("INSERT INTO test (x) VALUES", ((i,) for i in range(3))) + + res = client.execute("SELECT sum(x) FROM test WHERE x > %(minv)i", {"minv": 2}) + assert res[0][0] == 100 + + sentry_sdk.flush() + + spans = [item.payload for item in items] + assert len(spans) > 1 # sanity check that db spans were actually streamed + + for span in spans: + attribute_keys = { + attribute["name"] if isinstance(attribute, dict) else attribute + for attribute in span.get("attributes", {}) + } + assert "db.params" not in attribute_keys + assert "db.result" not in attribute_keys + + +def test_clickhouse_client_send_data_generator_with_data_collection( + sentry_init, capture_events +) -> None: + sentry_init( + integrations=[ClickhouseDriverIntegration()], + _experiments={"data_collection": {"database_query_data": True}}, + ) + events = capture_events() + + client = Client("localhost") + client.execute("DROP TABLE IF EXISTS test") + client.execute("CREATE TABLE test (x Int32) ENGINE = Memory") + client.execute("INSERT INTO test (x) VALUES", ((i,) for i in range(3))) + + res = client.execute("SELECT sum(x) FROM test") + assert res[0][0] == 3 + + capture_message("hi") + + (event,) = events + + (insert_breadcrumb,) = [ + crumb + for crumb in event["breadcrumbs"]["values"] + if crumb["message"] == "INSERT INTO test (x) VALUES" + ] + + assert insert_breadcrumb["data"]["db.params"] == [[0], [1], [2]] + + +def test_clickhouse_client_send_data_generator_with_data_collection_disabled( + sentry_init, capture_events +) -> None: + sentry_init( + integrations=[ClickhouseDriverIntegration()], + _experiments={"data_collection": {"database_query_data": False}}, + ) + events = capture_events() + + client = Client("localhost") + client.execute("DROP TABLE IF EXISTS test") + client.execute("CREATE TABLE test (x Int32) ENGINE = Memory") + client.execute("INSERT INTO test (x) VALUES", ((i,) for i in range(3))) + + res = client.execute("SELECT sum(x) FROM test") + assert res[0][0] == 3 + + capture_message("hi") + + (event,) = events + + (insert_breadcrumb,) = [ + crumb + for crumb in event["breadcrumbs"]["values"] + if crumb["message"] == "INSERT INTO test (x) VALUES" + ] + + assert "db.params" not in insert_breadcrumb["data"] + + @pytest.mark.parametrize("span_streaming", [True, False]) def test_clickhouse_client_spans( sentry_init,