From 9334dffbf99770d4c86e5d13d794c12bf557ec10 Mon Sep 17 00:00:00 2001 From: Ivana Kellyer Date: Wed, 23 Oct 2024 14:20:34 +0200 Subject: [PATCH 1/4] clickhouse fixes --- sentry_sdk/integrations/clickhouse_driver.py | 73 ++++++++++---------- 1 file changed, 36 insertions(+), 37 deletions(-) diff --git a/sentry_sdk/integrations/clickhouse_driver.py b/sentry_sdk/integrations/clickhouse_driver.py index 245ea0ef71..25f44559fc 100644 --- a/sentry_sdk/integrations/clickhouse_driver.py +++ b/sentry_sdk/integrations/clickhouse_driver.py @@ -9,7 +9,7 @@ ensure_integration_enabled, ) -from typing import TYPE_CHECKING, TypeVar +from typing import TYPE_CHECKING, Any, TypeVar # Hack to get new Python features working in older versions # without introducing a hard dependency on `typing_extensions` @@ -93,17 +93,17 @@ def _inner(*args: P.args, **kwargs: P.kwargs) -> T: connection._sentry_span = span # type: ignore[attr-defined] - _set_db_data(span, connection) - - if should_send_default_pii(): - span.set_attribute("db.query.text", query) + data = _get_db_data(connection) + data["db.query.text"] = query if query_id: - span.set_attribute("db.query_id", query_id) + data["db.query_id"] = query_id if params and should_send_default_pii(): - connection._sentry_db_params = params - span.set_attribute("db.params", _serialize_span_attribute(params)) + data["db.params"] = params + + connection._sentry_db_data = data # type: ignore[attr-defined] + _set_on_span(span, data) # run the original code ret = f(*args, **kwargs) @@ -116,28 +116,20 @@ 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] + connection = args[0].connection + span = getattr(connection, "_sentry_span", None) # type: ignore[attr-defined] if span is not None: + data = getattr(connection, "_sentry_db_data", {}) + if res is not None and should_send_default_pii(): - span.set_attribute("db.result", _serialize_span_attribute(res)) + serialized_result = _serialize_span_attribute(res) + data["db.result"] = serialized_result + span.set_attribute("db.result", serialized_result) with capture_internal_exceptions(): - query = span.get_attribute("db.query.text") + query = data.pop("db.query.text", None) if query: - data = {} - for attr in ( - "db.params", - "db.result", - SPANDATA.DB_SYSTEM, - SPANDATA.DB_USER, - SPANDATA.SERVER_ADDRESS, - SPANDATA.SERVER_PORT, - ): - if span.get_attribute(attr): - data[attr] = span.get_attribute(attr) - sentry_sdk.add_breadcrumb( message=query, category="query", data=data ) @@ -152,20 +144,22 @@ def _inner_end(*args: P.args, **kwargs: P.kwargs) -> T: def _wrap_send_data(f: Callable[P, T]) -> Callable[P, T]: def _inner_send_data(*args: P.args, **kwargs: P.kwargs) -> T: instance = args[0] # type: clickhouse_driver.client.Client - data = args[2] + db_params_data = args[2] span = getattr(instance.connection, "_sentry_span", None) if span is not None: - _set_db_data(span, instance.connection) + data = _get_db_data(instance.connection) + _set_on_span(span, data) if should_send_default_pii(): db_params = ( - getattr(instance.connection, "_sentry_db_params", None) or [] + getattr(instance.connection, "_sentry_db_data", {}).get("db.params") + or [] ) - db_params.extend(data) + db_params.extend(db_params_data) span.set_attribute("db.params", _serialize_span_attribute(db_params)) try: - del instance.connection._sentry_db_params + del instance.connection._sentry_db_data except AttributeError: pass @@ -174,11 +168,16 @@ def _inner_send_data(*args: P.args, **kwargs: P.kwargs) -> T: return _inner_send_data -def _set_db_data( - span: Span, connection: clickhouse_driver.connection.Connection -) -> None: - span.set_attribute(SPANDATA.DB_SYSTEM, "clickhouse") - span.set_attribute(SPANDATA.SERVER_ADDRESS, connection.host) - span.set_attribute(SPANDATA.SERVER_PORT, connection.port) - span.set_attribute(SPANDATA.DB_NAME, connection.database) - span.set_attribute(SPANDATA.DB_USER, connection.user) +def _get_db_data(connection: clickhouse_driver.connection.Connection) -> dict[str, str]: + return { + SPANDATA.DB_SYSTEM: "clickhouse", + SPANDATA.SERVER_ADDRESS: connection.host, + SPANDATA.SERVER_PORT: connection.port, + SPANDATA.DB_NAME: connection.database, + SPANDATA.DB_USER: connection.user, + } + + +def _set_on_span(span: Span, data: dict[str, Any]): + for key, value in data.items(): + span.set_attribute(key, _serialize_span_attribute(value)) From b33f83cd405a379a3a2abd54037528645d507859 Mon Sep 17 00:00:00 2001 From: Ivana Kellyer Date: Wed, 23 Oct 2024 15:05:59 +0200 Subject: [PATCH 2/4] fix --- sentry_sdk/integrations/clickhouse_driver.py | 23 ++++++++++--------- .../test_clickhouse_driver.py | 18 +++++++-------- 2 files changed, 20 insertions(+), 21 deletions(-) diff --git a/sentry_sdk/integrations/clickhouse_driver.py b/sentry_sdk/integrations/clickhouse_driver.py index 25f44559fc..d6d720590b 100644 --- a/sentry_sdk/integrations/clickhouse_driver.py +++ b/sentry_sdk/integrations/clickhouse_driver.py @@ -136,6 +136,12 @@ def _inner_end(*args: P.args, **kwargs: P.kwargs) -> T: span.finish() + try: + del connection._sentry_db_data + del connection._sentry_span + except AttributeError: + pass + return res return _inner_end @@ -143,25 +149,20 @@ def _inner_end(*args: P.args, **kwargs: P.kwargs) -> T: def _wrap_send_data(f: Callable[P, T]) -> Callable[P, T]: def _inner_send_data(*args: P.args, **kwargs: P.kwargs) -> T: - instance = args[0] # type: clickhouse_driver.client.Client + connection = args[0].connection db_params_data = args[2] - span = getattr(instance.connection, "_sentry_span", None) + span = getattr(connection, "_sentry_span", None) if span is not None: - data = _get_db_data(instance.connection) + data = _get_db_data(connection) _set_on_span(span, data) if should_send_default_pii(): - db_params = ( - getattr(instance.connection, "_sentry_db_data", {}).get("db.params") - or [] - ) + saved_db_data = getattr(connection, "_sentry_db_data", {}) + db_params = saved_db_data.get("db.params") or [] db_params.extend(db_params_data) + saved_db_data["db.params"] = db_params span.set_attribute("db.params", _serialize_span_attribute(db_params)) - try: - del instance.connection._sentry_db_data - except AttributeError: - pass return f(*args, **kwargs) diff --git a/tests/integrations/clickhouse_driver/test_clickhouse_driver.py b/tests/integrations/clickhouse_driver/test_clickhouse_driver.py index 5da77ce13d..8b129abe77 100644 --- a/tests/integrations/clickhouse_driver/test_clickhouse_driver.py +++ b/tests/integrations/clickhouse_driver/test_clickhouse_driver.py @@ -168,7 +168,7 @@ def test_clickhouse_client_breadcrumbs_with_pii(sentry_init, capture_events) -> "db.user": "default", "server.address": "localhost", "server.port": 9000, - "db.params": '[{"x": 100}]', + "db.params": [{"x": 100}], }, "message": "INSERT INTO test (x) VALUES", "type": "default", @@ -181,7 +181,7 @@ def test_clickhouse_client_breadcrumbs_with_pii(sentry_init, capture_events) -> "db.user": "default", "server.address": "localhost", "server.port": 9000, - "db.params": "[[170], [200]]", + "db.params": [[170], [200]], }, "message": "INSERT INTO test (x) VALUES", "type": "default", @@ -195,7 +195,7 @@ def test_clickhouse_client_breadcrumbs_with_pii(sentry_init, capture_events) -> "server.address": "localhost", "server.port": 9000, "db.result": "[[370]]", - "db.params": '{"minv": 150}', + "db.params": {"minv": 150}, }, "message": "SELECT sum(x) FROM test WHERE x > 150", "type": "default", @@ -348,9 +348,7 @@ def test_clickhouse_client_spans( assert event["spans"] == expected_spans -def test_clickhouse_client_spans_with_pii( - sentry_init, capture_events, capture_envelopes -) -> None: +def test_clickhouse_client_spans_with_pii(sentry_init, capture_events) -> None: sentry_init( integrations=[ClickhouseDriverIntegration()], _experiments={"record_sql_params": True}, @@ -646,7 +644,7 @@ def test_clickhouse_dbapi_breadcrumbs_with_pii(sentry_init, capture_events) -> N "db.user": "default", "server.address": "localhost", "server.port": 9000, - "db.params": '[{"x": 100}]', + "db.params": [{"x": 100}], }, "message": "INSERT INTO test (x) VALUES", "type": "default", @@ -659,7 +657,7 @@ def test_clickhouse_dbapi_breadcrumbs_with_pii(sentry_init, capture_events) -> N "db.user": "default", "server.address": "localhost", "server.port": 9000, - "db.params": "[[170], [200]]", + "db.params": [[170], [200]], }, "message": "INSERT INTO test (x) VALUES", "type": "default", @@ -672,8 +670,8 @@ def test_clickhouse_dbapi_breadcrumbs_with_pii(sentry_init, capture_events) -> N "db.user": "default", "server.address": "localhost", "server.port": 9000, - "db.params": '{"minv": 150}', - "db.result": '[[["370"]], [["\'sum(x)\'", "\'Int64\'"]]]', + "db.params": {"minv": 150}, + "db.result": '[[[370]], [["sum(x)", "Int64"]]]', }, "message": "SELECT sum(x) FROM test WHERE x > 150", "type": "default", From 36c1ebaf88a193cff56eec97ccab9f5c7564dc07 Mon Sep 17 00:00:00 2001 From: Ivana Kellyer Date: Wed, 23 Oct 2024 15:12:03 +0200 Subject: [PATCH 3/4] consistent serialization --- sentry_sdk/integrations/clickhouse_driver.py | 5 ++--- .../clickhouse_driver/test_clickhouse_driver.py | 8 ++++---- 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/sentry_sdk/integrations/clickhouse_driver.py b/sentry_sdk/integrations/clickhouse_driver.py index d6d720590b..7bd7b59524 100644 --- a/sentry_sdk/integrations/clickhouse_driver.py +++ b/sentry_sdk/integrations/clickhouse_driver.py @@ -123,9 +123,8 @@ def _inner_end(*args: P.args, **kwargs: P.kwargs) -> T: data = getattr(connection, "_sentry_db_data", {}) if res is not None and should_send_default_pii(): - serialized_result = _serialize_span_attribute(res) - data["db.result"] = serialized_result - span.set_attribute("db.result", serialized_result) + data["db.result"] = res + span.set_attribute("db.result", _serialize_span_attribute(res)) with capture_internal_exceptions(): query = data.pop("db.query.text", None) diff --git a/tests/integrations/clickhouse_driver/test_clickhouse_driver.py b/tests/integrations/clickhouse_driver/test_clickhouse_driver.py index 8b129abe77..381cbaafd1 100644 --- a/tests/integrations/clickhouse_driver/test_clickhouse_driver.py +++ b/tests/integrations/clickhouse_driver/test_clickhouse_driver.py @@ -194,7 +194,7 @@ def test_clickhouse_client_breadcrumbs_with_pii(sentry_init, capture_events) -> "db.user": "default", "server.address": "localhost", "server.port": 9000, - "db.result": "[[370]]", + "db.result": [[370]], "db.params": {"minv": 150}, }, "message": "SELECT sum(x) FROM test WHERE x > 150", @@ -618,7 +618,7 @@ def test_clickhouse_dbapi_breadcrumbs_with_pii(sentry_init, capture_events) -> N "db.user": "default", "server.address": "localhost", "server.port": 9000, - "db.result": "[[], []]", + "db.result": [[], []], }, "message": "DROP TABLE IF EXISTS test", "type": "default", @@ -631,7 +631,7 @@ def test_clickhouse_dbapi_breadcrumbs_with_pii(sentry_init, capture_events) -> N "db.user": "default", "server.address": "localhost", "server.port": 9000, - "db.result": "[[], []]", + "db.result": [[], []], }, "message": "CREATE TABLE test (x Int32) ENGINE = Memory", "type": "default", @@ -671,7 +671,7 @@ def test_clickhouse_dbapi_breadcrumbs_with_pii(sentry_init, capture_events) -> N "server.address": "localhost", "server.port": 9000, "db.params": {"minv": 150}, - "db.result": '[[[370]], [["sum(x)", "Int64"]]]', + "db.result": [[["370"]], [["'sum(x)'", "'Int64'"]]], }, "message": "SELECT sum(x) FROM test WHERE x > 150", "type": "default", From 63293bfe01d32fad24ebd097a24f5907478506f9 Mon Sep 17 00:00:00 2001 From: Ivana Kellyer Date: Tue, 19 Nov 2024 17:09:34 +0100 Subject: [PATCH 4/4] old py compat --- sentry_sdk/integrations/clickhouse_driver.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/sentry_sdk/integrations/clickhouse_driver.py b/sentry_sdk/integrations/clickhouse_driver.py index 3efe91471f..0c83a30824 100644 --- a/sentry_sdk/integrations/clickhouse_driver.py +++ b/sentry_sdk/integrations/clickhouse_driver.py @@ -9,7 +9,7 @@ ensure_integration_enabled, ) -from typing import TYPE_CHECKING, Any, TypeVar +from typing import TYPE_CHECKING, Any, Dict, TypeVar # Hack to get new Python features working in older versions # without introducing a hard dependency on `typing_extensions` @@ -169,7 +169,7 @@ def _inner_send_data(*args: P.args, **kwargs: P.kwargs) -> T: return _inner_send_data -def _get_db_data(connection: clickhouse_driver.connection.Connection) -> dict[str, str]: +def _get_db_data(connection: clickhouse_driver.connection.Connection) -> Dict[str, str]: return { SPANDATA.DB_SYSTEM: "clickhouse", SPANDATA.SERVER_ADDRESS: connection.host, @@ -179,6 +179,6 @@ def _get_db_data(connection: clickhouse_driver.connection.Connection) -> dict[st } -def _set_on_span(span: Span, data: dict[str, Any]): +def _set_on_span(span: Span, data: Dict[str, Any]): for key, value in data.items(): span.set_attribute(key, _serialize_span_attribute(value))