From 07c8f498d45c59e0db1e81e319b885e5dd257e8e Mon Sep 17 00:00:00 2001 From: Marcel Coetzee Date: Mon, 1 Jul 2024 15:39:25 +0200 Subject: [PATCH] Refactor ClickHouse deployment type detection Signed-off-by: Marcel Coetzee --- dlt/destinations/impl/clickhouse/sql_client.py | 6 ------ tests/load/clickhouse/test_clickhouse_adapter.py | 3 ++- tests/load/clickhouse/utils.py | 8 ++++++++ 3 files changed, 10 insertions(+), 7 deletions(-) create mode 100644 tests/load/clickhouse/utils.py diff --git a/dlt/destinations/impl/clickhouse/sql_client.py b/dlt/destinations/impl/clickhouse/sql_client.py index 9b1617f043..8d42bbef58 100644 --- a/dlt/destinations/impl/clickhouse/sql_client.py +++ b/dlt/destinations/impl/clickhouse/sql_client.py @@ -217,9 +217,3 @@ def _make_database_exception(cls, ex: Exception) -> Exception: @staticmethod def is_dbapi_exception(ex: Exception) -> bool: return isinstance(ex, clickhouse_driver.dbapi.Error) - - def _get_deployment_type(self) -> TDeployment: - cloud_mode = int(self.execute_sql(""" - SELECT value FROM system.settings WHERE name = 'cloud_mode' - """)[0][0]) - return "ClickHouseCloud" if cloud_mode else "ClickHouseOSS" diff --git a/tests/load/clickhouse/test_clickhouse_adapter.py b/tests/load/clickhouse/test_clickhouse_adapter.py index 88a21f9bc8..4abf094576 100644 --- a/tests/load/clickhouse/test_clickhouse_adapter.py +++ b/tests/load/clickhouse/test_clickhouse_adapter.py @@ -3,6 +3,7 @@ import dlt from dlt.destinations.adapters import clickhouse_adapter from dlt.destinations.impl.clickhouse.sql_client import TDeployment +from tests.load.clickhouse.utils import get_deployment_type from tests.pipeline.utils import assert_load_info @@ -34,7 +35,7 @@ def not_annotated_resource() -> Generator[Dict[str, int], None, None]: pipe = dlt.pipeline(pipeline_name="adapter_test", destination="clickhouse", dev_mode=True) with pipe.sql_client() as client: - deployment_type: TDeployment = client._get_deployment_type() + deployment_type: TDeployment = get_deployment_type(client) if deployment_type == "ClickHouseCloud": pack = pipe.run( diff --git a/tests/load/clickhouse/utils.py b/tests/load/clickhouse/utils.py new file mode 100644 index 0000000000..809e929261 --- /dev/null +++ b/tests/load/clickhouse/utils.py @@ -0,0 +1,8 @@ +from dlt.destinations.impl.clickhouse.sql_client import TDeployment, ClickHouseSqlClient + + +def get_deployment_type(client: ClickHouseSqlClient) -> TDeployment: + cloud_mode = int(client.execute_sql(""" + SELECT value FROM system.settings WHERE name = 'cloud_mode' + """)[0][0]) + return "ClickHouseCloud" if cloud_mode else "ClickHouseOSS"