diff --git a/sentry_sdk/integrations/clickhouse_driver.py b/sentry_sdk/integrations/clickhouse_driver.py index 4e9901acaf..b2e440a50f 100644 --- a/sentry_sdk/integrations/clickhouse_driver.py +++ b/sentry_sdk/integrations/clickhouse_driver.py @@ -15,7 +15,7 @@ # from: https://stackoverflow.com/a/71944042/300572 if TYPE_CHECKING: from collections.abc import Iterator - from typing import Any, Callable, ParamSpec, Union + from typing import Any, Callable, Optional, ParamSpec, Union else: # Fake ParamSpec class ParamSpec: @@ -79,7 +79,7 @@ def _inner(*args: "P.args", **kwargs: "P.kwargs") -> "T": if client.get_integration(ClickhouseDriverIntegration) is None: return f(*args, **kwargs) - connection = args[0] + connection: "Connection" = args[0] query = args[1] query_id = args[2] if len(args) > 2 else kwargs.get("query_id") params = args[3] if len(args) > 3 else kwargs.get("params") @@ -95,6 +95,16 @@ def _inner(*args: "P.args", **kwargs: "P.kwargs") -> "T": SPANDATA.DB_QUERY_TEXT: str(query), }, ) + + connection._query = query + connection._breadcrumb_data = { + SPANDATA.DB_SYSTEM: "clickhouse", + SPANDATA.DB_NAME: connection.database, + SPANDATA.DB_DRIVER_NAME: "clickhouse-driver", + SPANDATA.SERVER_ADDRESS: connection.host, + SPANDATA.SERVER_PORT: connection.port, + SPANDATA.DB_USER: connection.user, + } else: span = sentry_sdk.start_span( op=OP.DB, @@ -114,7 +124,7 @@ def _inner(*args: "P.args", **kwargs: "P.kwargs") -> "T": elif should_send_default_pii(): span.set_data("db.params", params) - connection._sentry_span = span # type: ignore[attr-defined] + connection._sentry_span = span if span is not None: _set_db_data(span, connection) @@ -130,8 +140,31 @@ def _inner(*args: "P.args", **kwargs: "P.kwargs") -> "T": def _wrap_end(f: "Callable[P, T]") -> "Callable[P, T]": def _inner_end(*args: "P.args", **kwargs: "P.kwargs") -> "T": res = f(*args, **kwargs) - instance = args[0] - span = getattr(instance.connection, "_sentry_span", None) # type: ignore[attr-defined] + instance: "Client" = args[0] + + query = getattr(instance.connection, "_query", None) + breadcrumb_data: "Optional[dict[str, Any]]" = getattr( + instance.connection, "_breadcrumb_data", None + ) + + if query is not None and breadcrumb_data is not None: + client_options = sentry_sdk.get_client().options + if ( + has_data_collection_enabled(client_options) + and client_options["data_collection"]["database_query_data"] + ) or ( + not has_data_collection_enabled(client_options) + and should_send_default_pii() + ): + breadcrumb_data = {"db.result": res, **breadcrumb_data} + + sentry_sdk.get_isolation_scope().add_breadcrumb( + message=query, + category="query", + data={"db.result": res, **breadcrumb_data}, + ) + + span = getattr(instance.connection, "_sentry_span", None) if span is None: return res diff --git a/tests/integrations/clickhouse_driver/test_clickhouse_driver.py b/tests/integrations/clickhouse_driver/test_clickhouse_driver.py index 1cec232548..b6d45a2365 100644 --- a/tests/integrations/clickhouse_driver/test_clickhouse_driver.py +++ b/tests/integrations/clickhouse_driver/test_clickhouse_driver.py @@ -21,9 +21,13 @@ EXPECT_PARAMS_IN_SELECT = False -def test_clickhouse_client_breadcrumbs(sentry_init, capture_events) -> None: +@pytest.mark.parametrize("span_streaming", [True, False]) +def test_clickhouse_client_breadcrumbs( + sentry_init, capture_events, span_streaming +) -> None: sentry_init( integrations=[ClickhouseDriverIntegration()], + trace_lifecycle="stream" if span_streaming else "static", _experiments={"record_sql_params": True}, ) events = capture_events() @@ -127,6 +131,112 @@ def test_clickhouse_client_breadcrumbs(sentry_init, capture_events) -> None: assert actual_query_breadcrumbs == expected_breadcrumbs +def test_clickhouse_client_breadcrumbs_span_streaming( + sentry_init, capture_events +) -> None: + sentry_init( + integrations=[ClickhouseDriverIntegration()], + trace_lifecycle="stream", + _experiments={"record_sql_params": 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.driver.name": "clickhouse-driver", + "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.driver.name": "clickhouse-driver", + "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.driver.name": "clickhouse-driver", + "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.driver.name": "clickhouse-driver", + "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.driver.name": "clickhouse-driver", + "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) + + actual_query_breadcrumbs = [ + breadcrumb + for breadcrumb in event["breadcrumbs"]["values"] + if breadcrumb["category"] == "query" + ] + + assert actual_query_breadcrumbs == expected_breadcrumbs + + def test_clickhouse_client_breadcrumbs_with_pii(sentry_init, capture_events) -> None: sentry_init( integrations=[ClickhouseDriverIntegration()], @@ -229,6 +339,107 @@ def test_clickhouse_client_breadcrumbs_with_pii(sentry_init, capture_events) -> assert event["breadcrumbs"]["values"] == expected_breadcrumbs +def test_clickhouse_client_breadcrumbs_with_pii_span_streaming( + sentry_init, capture_events +) -> None: + sentry_init( + integrations=[ClickhouseDriverIntegration()], + send_default_pii=True, + trace_lifecycle="stream", + _experiments={"record_sql_params": 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.driver.name": "clickhouse-driver", + "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.driver.name": "clickhouse-driver", + "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.driver.name": "clickhouse-driver", + "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.driver.name": "clickhouse-driver", + "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.driver.name": "clickhouse-driver", + "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 + + def test_clickhouse_client_breadcrumbs_with_data_collection( sentry_init, capture_events ) -> None: @@ -332,12 +543,13 @@ def test_clickhouse_client_breadcrumbs_with_data_collection( assert event["breadcrumbs"]["values"] == expected_breadcrumbs -def test_clickhouse_client_breadcrumbs_with_data_collection_disabled( +def test_clickhouse_client_breadcrumbs_with_data_collection_span_streaming( sentry_init, capture_events ) -> None: sentry_init( integrations=[ClickhouseDriverIntegration()], - _experiments={"data_collection": {"database_query_data": False}}, + trace_lifecycle="stream", + _experiments={"data_collection": {"database_query_data": True}}, ) events = capture_events() @@ -359,6 +571,7 @@ def test_clickhouse_client_breadcrumbs_with_data_collection_disabled( "category": "query", "data": { "db.system": "clickhouse", + "db.driver.name": "clickhouse-driver", "db.name": "", "db.user": "default", "server.address": "localhost", @@ -371,6 +584,7 @@ def test_clickhouse_client_breadcrumbs_with_data_collection_disabled( "category": "query", "data": { "db.system": "clickhouse", + "db.driver.name": "clickhouse-driver", "db.name": "", "db.user": "default", "server.address": "localhost", @@ -383,6 +597,7 @@ def test_clickhouse_client_breadcrumbs_with_data_collection_disabled( "category": "query", "data": { "db.system": "clickhouse", + "db.driver.name": "clickhouse-driver", "db.name": "", "db.user": "default", "server.address": "localhost", @@ -395,6 +610,7 @@ def test_clickhouse_client_breadcrumbs_with_data_collection_disabled( "category": "query", "data": { "db.system": "clickhouse", + "db.driver.name": "clickhouse-driver", "db.name": "", "db.user": "default", "server.address": "localhost", @@ -407,6 +623,7 @@ def test_clickhouse_client_breadcrumbs_with_data_collection_disabled( "category": "query", "data": { "db.system": "clickhouse", + "db.driver.name": "clickhouse-driver", "db.name": "", "db.user": "default", "server.address": "localhost", @@ -425,18 +642,12 @@ def test_clickhouse_client_breadcrumbs_with_data_collection_disabled( 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( +def test_clickhouse_client_breadcrumbs_with_data_collection_disabled( sentry_init, capture_events ) -> None: sentry_init( integrations=[ClickhouseDriverIntegration()], - send_default_pii=True, _experiments={"data_collection": {"database_query_data": False}}, ) events = capture_events() @@ -531,12 +742,13 @@ def test_clickhouse_client_breadcrumbs_data_collection_overrides_pii( assert "db.result" not in crumb["data"] -def test_clickhouse_client_breadcrumbs_with_data_collection_default( +def test_clickhouse_client_breadcrumbs_with_data_collection_disabled_span_streaming( sentry_init, capture_events ) -> None: sentry_init( integrations=[ClickhouseDriverIntegration()], - _experiments={"data_collection": {}}, + trace_lifecycle="stream", + _experiments={"data_collection": {"database_query_data": False}}, ) events = capture_events() @@ -558,11 +770,11 @@ def test_clickhouse_client_breadcrumbs_with_data_collection_default( "category": "query", "data": { "db.system": "clickhouse", + "db.driver.name": "clickhouse-driver", "db.name": "", "db.user": "default", "server.address": "localhost", "server.port": 9000, - "db.result": [], }, "message": "DROP TABLE IF EXISTS test", "type": "default", @@ -571,11 +783,11 @@ def test_clickhouse_client_breadcrumbs_with_data_collection_default( "category": "query", "data": { "db.system": "clickhouse", + "db.driver.name": "clickhouse-driver", "db.name": "", "db.user": "default", "server.address": "localhost", "server.port": 9000, - "db.result": [], }, "message": "CREATE TABLE test (x Int32) ENGINE = Memory", "type": "default", @@ -584,11 +796,11 @@ def test_clickhouse_client_breadcrumbs_with_data_collection_default( "category": "query", "data": { "db.system": "clickhouse", + "db.driver.name": "clickhouse-driver", "db.name": "", "db.user": "default", "server.address": "localhost", "server.port": 9000, - "db.params": [{"x": 100}], }, "message": "INSERT INTO test (x) VALUES", "type": "default", @@ -597,11 +809,11 @@ def test_clickhouse_client_breadcrumbs_with_data_collection_default( "category": "query", "data": { "db.system": "clickhouse", + "db.driver.name": "clickhouse-driver", "db.name": "", "db.user": "default", "server.address": "localhost", "server.port": 9000, - "db.params": [[170], [200]], }, "message": "INSERT INTO test (x) VALUES", "type": "default", @@ -610,74 +822,535 @@ def test_clickhouse_client_breadcrumbs_with_data_collection_default( "category": "query", "data": { "db.system": "clickhouse", + "db.driver.name": "clickhouse-driver", "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 + assert event["breadcrumbs"]["values"] == expected_breadcrumbs + + +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_data_collection_overrides_pii_span_streaming( + sentry_init, capture_events +) -> None: + sentry_init( + integrations=[ClickhouseDriverIntegration()], + send_default_pii=True, + trace_lifecycle="stream", + _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.driver.name": "clickhouse-driver", + "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.driver.name": "clickhouse-driver", + "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.driver.name": "clickhouse-driver", + "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.driver.name": "clickhouse-driver", + "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.driver.name": "clickhouse-driver", + "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 + + +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_breadcrumbs_with_data_collection_default_span_streaming( + sentry_init, capture_events +) -> None: + sentry_init( + integrations=[ClickhouseDriverIntegration()], + trace_lifecycle="stream", + _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.driver.name": "clickhouse-driver", + "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.driver.name": "clickhouse-driver", + "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.driver.name": "clickhouse-driver", + "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.driver.name": "clickhouse-driver", + "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.driver.name": "clickhouse-driver", + "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 + + +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, + trace_lifecycle="stream", + _experiments={ + "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_span_streaming_with_data_collection( - sentry_init, capture_items +def test_clickhouse_client_send_data_generator_with_data_collection_disabled( + sentry_init, capture_events ) -> None: - # Streamed spans never carry db.params/db.result, regardless of the - # data_collection configuration sentry_init( integrations=[ClickhouseDriverIntegration()], - traces_sample_rate=1.0, - trace_lifecycle="stream", - _experiments={ - "data_collection": {"database_query_data": True}, - }, + _experiments={"data_collection": {"database_query_data": False}}, ) - items = capture_items("span") + events = capture_events() - 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))) + 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 WHERE x > %(minv)i", {"minv": 2}) - assert res[0][0] == 100 + res = client.execute("SELECT sum(x) FROM test") + assert res[0][0] == 3 - sentry_sdk.flush() + capture_message("hi") - spans = [item.payload for item in items] - assert len(spans) > 1 # sanity check that db spans were actually streamed + (event,) = events - 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 + (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"] -def test_clickhouse_client_send_data_generator_with_data_collection( +def test_clickhouse_client_send_data_generator_span_streaming_with_data_collection( sentry_init, capture_events ) -> None: sentry_init( integrations=[ClickhouseDriverIntegration()], + trace_lifecycle="stream", _experiments={"data_collection": {"database_query_data": True}}, ) events = capture_events() @@ -700,14 +1373,15 @@ def test_clickhouse_client_send_data_generator_with_data_collection( if crumb["message"] == "INSERT INTO test (x) VALUES" ] - assert insert_breadcrumb["data"]["db.params"] == [[0], [1], [2]] + assert "db.params" not in insert_breadcrumb["data"] -def test_clickhouse_client_send_data_generator_with_data_collection_disabled( +def test_clickhouse_client_send_data_generator_span_streaming_with_data_collection_disabled( sentry_init, capture_events ) -> None: sentry_init( integrations=[ClickhouseDriverIntegration()], + trace_lifecycle="stream", _experiments={"data_collection": {"database_query_data": False}}, ) events = capture_events() @@ -1376,6 +2050,107 @@ def test_clickhouse_dbapi_breadcrumbs(sentry_init, capture_events) -> None: assert event["breadcrumbs"]["values"] == expected_breadcrumbs +def test_clickhouse_dbapi_breadcrumbs_span_streaming( + sentry_init, capture_events +) -> None: + sentry_init( + integrations=[ClickhouseDriverIntegration()], + trace_lifecycle="stream", + ) + events = capture_events() + + conn = connect("clickhouse://localhost") + cursor = conn.cursor() + cursor.execute("DROP TABLE IF EXISTS test") + cursor.execute("CREATE TABLE test (x Int32) ENGINE = Memory") + cursor.executemany("INSERT INTO test (x) VALUES", [{"x": 100}]) + cursor.executemany("INSERT INTO test (x) VALUES", [[170], [200]]) + cursor.execute("SELECT sum(x) FROM test WHERE x > %(minv)i", {"minv": 150}) + res = cursor.fetchall() + + assert res[0][0] == 370 + + capture_message("hi") + + (event,) = events + + expected_breadcrumbs = [ + { + "category": "query", + "data": { + "db.system": "clickhouse", + "db.driver.name": "clickhouse-driver", + "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.driver.name": "clickhouse-driver", + "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.driver.name": "clickhouse-driver", + "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.driver.name": "clickhouse-driver", + "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.driver.name": "clickhouse-driver", + "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 + + def test_clickhouse_dbapi_breadcrumbs_with_pii(sentry_init, capture_events) -> None: sentry_init( integrations=[ClickhouseDriverIntegration()], @@ -1479,6 +2254,108 @@ def test_clickhouse_dbapi_breadcrumbs_with_pii(sentry_init, capture_events) -> N assert event["breadcrumbs"]["values"] == expected_breadcrumbs +def test_clickhouse_dbapi_breadcrumbs_with_pii_span_streaming( + sentry_init, capture_events +) -> None: + sentry_init( + integrations=[ClickhouseDriverIntegration()], + send_default_pii=True, + trace_lifecycle="stream", + ) + events = capture_events() + + conn = connect("clickhouse://localhost") + cursor = conn.cursor() + cursor.execute("DROP TABLE IF EXISTS test") + cursor.execute("CREATE TABLE test (x Int32) ENGINE = Memory") + cursor.executemany("INSERT INTO test (x) VALUES", [{"x": 100}]) + cursor.executemany("INSERT INTO test (x) VALUES", [[170], [200]]) + cursor.execute("SELECT sum(x) FROM test WHERE x > %(minv)i", {"minv": 150}) + res = cursor.fetchall() + + assert res[0][0] == 370 + + capture_message("hi") + + (event,) = events + + expected_breadcrumbs = [ + { + "category": "query", + "data": { + "db.system": "clickhouse", + "db.driver.name": "clickhouse-driver", + "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.driver.name": "clickhouse-driver", + "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.driver.name": "clickhouse-driver", + "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.driver.name": "clickhouse-driver", + "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.driver.name": "clickhouse-driver", + "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 + + @pytest.mark.parametrize("span_streaming", [True, False]) def test_clickhouse_dbapi_spans( sentry_init,