diff --git a/sentry_sdk/integrations/clickhouse_driver.py b/sentry_sdk/integrations/clickhouse_driver.py index b4cc2860e7..7bbea94210 100644 --- a/sentry_sdk/integrations/clickhouse_driver.py +++ b/sentry_sdk/integrations/clickhouse_driver.py @@ -30,7 +30,9 @@ def __getitem__(self, _): try: - import clickhouse_driver # type: ignore[import-not-found] + from clickhouse_driver import VERSION # type: ignore[import-not-found] + from clickhouse_driver.client import Client # type: ignore[import-not-found] + from clickhouse_driver.connection import Connection # type: ignore[import-not-found] except ImportError: raise DidNotEnable("clickhouse-driver not installed.") @@ -42,29 +44,23 @@ class ClickhouseDriverIntegration(Integration): @staticmethod def setup_once() -> None: - _check_minimum_version(ClickhouseDriverIntegration, clickhouse_driver.VERSION) + _check_minimum_version(ClickhouseDriverIntegration, VERSION) # Every query is done using the Connection's `send_query` function - clickhouse_driver.connection.Connection.send_query = _wrap_start( - clickhouse_driver.connection.Connection.send_query - ) + Connection.send_query = _wrap_start(Connection.send_query) # If the query contains parameters then the send_data function is used to send those parameters to clickhouse _wrap_send_data() # Every query ends either with the Client's `receive_end_of_query` (no result expected) # or its `receive_result` (result expected) - clickhouse_driver.client.Client.receive_end_of_query = _wrap_end( - clickhouse_driver.client.Client.receive_end_of_query - ) - if hasattr(clickhouse_driver.client.Client, "receive_end_of_insert_query"): + Client.receive_end_of_query = _wrap_end(Client.receive_end_of_query) + if hasattr(Client, "receive_end_of_insert_query"): # In 0.2.7, insert queries are handled separately via `receive_end_of_insert_query` - clickhouse_driver.client.Client.receive_end_of_insert_query = _wrap_end( - clickhouse_driver.client.Client.receive_end_of_insert_query + Client.receive_end_of_insert_query = _wrap_end( + Client.receive_end_of_insert_query ) - clickhouse_driver.client.Client.receive_result = _wrap_end( - clickhouse_driver.client.Client.receive_result - ) + Client.receive_result = _wrap_end(Client.receive_result) P = ParamSpec("P") @@ -128,7 +124,7 @@ def _inner_end(*args: "P.args", **kwargs: "P.kwargs") -> "T": def _wrap_send_data() -> None: - original_send_data = clickhouse_driver.client.Client.send_data + original_send_data = Client.send_data def _inner_send_data( # type: ignore[no-untyped-def] # clickhouse-driver does not type send_data self, sample_block, data, types_check=False, columnar=False, *args, **kwargs @@ -164,12 +160,10 @@ def wrapped_generator() -> "Iterator[Any]": self, sample_block, data, types_check, columnar, *args, **kwargs ) - clickhouse_driver.client.Client.send_data = _inner_send_data + Client.send_data = _inner_send_data -def _set_db_data( - span: "Span", connection: "clickhouse_driver.connection.Connection" -) -> None: +def _set_db_data(span: "Span", connection: "Connection") -> None: span.set_data(SPANDATA.DB_SYSTEM, "clickhouse") span.set_data(SPANDATA.SERVER_ADDRESS, connection.host) span.set_data(SPANDATA.SERVER_PORT, connection.port) diff --git a/tests/test_shadowed_module.py b/tests/test_shadowed_module.py index e1171dd103..e146ab2f4f 100644 --- a/tests/test_shadowed_module.py +++ b/tests/test_shadowed_module.py @@ -29,7 +29,6 @@ def pytest_generate_tests(metafunc): # Temporarily skip some integrations submodule_names - { - "clickhouse_driver", "grpc", "litellm", "opentelemetry",