From 027d52a726a958ef2c8af908e0772afcc391ee60 Mon Sep 17 00:00:00 2001
From: Marcel Coetzee <marcel@mooncoon.com>
Date: Tue, 16 Apr 2024 23:09:49 +0200
Subject: [PATCH] Add test for clickhouse config settings #1055

Signed-off-by: Marcel Coetzee <marcel@mooncoon.com>
---
 .../impl/clickhouse/clickhouse.py             |  1 +
 .../impl/clickhouse/sql_client.py             | 11 +-----
 .../test_clickhouse_configuration.py          | 38 ++++++++++++++++++-
 3 files changed, 39 insertions(+), 11 deletions(-)

diff --git a/dlt/destinations/impl/clickhouse/clickhouse.py b/dlt/destinations/impl/clickhouse/clickhouse.py
index 3489c5cacb..8575c5a52c 100644
--- a/dlt/destinations/impl/clickhouse/clickhouse.py
+++ b/dlt/destinations/impl/clickhouse/clickhouse.py
@@ -254,6 +254,7 @@ def __init__(
                         settings={
                             "allow_experimental_lightweight_delete": 1,
                             "allow_experimental_object_type": 1,
+                            "enable_http_compression": 1,
                         },
                         compression=compression,
                     )
diff --git a/dlt/destinations/impl/clickhouse/sql_client.py b/dlt/destinations/impl/clickhouse/sql_client.py
index 5710e41e49..8fafa7ba74 100644
--- a/dlt/destinations/impl/clickhouse/sql_client.py
+++ b/dlt/destinations/impl/clickhouse/sql_client.py
@@ -14,7 +14,6 @@
 from clickhouse_driver.dbapi import OperationalError  # type: ignore[import-untyped]
 from clickhouse_driver.dbapi.extras import DictCursor  # type: ignore[import-untyped]
 
-import dlt
 from dlt.common.destination import DestinationCapabilitiesContext
 from dlt.destinations.exceptions import (
     DatabaseUndefinedRelation,
@@ -128,7 +127,7 @@ def _list_tables(self) -> List[str]:
     def execute_query(
         self, query: AnyStr, *args: Any, **kwargs: Any
     ) -> Iterator[ClickHouseDBApiCursorImpl]:
-        assert isinstance(query, str), "Query must be a string"
+        assert isinstance(query, str), "Query must be a string."
 
         db_args = kwargs.copy()
 
@@ -136,14 +135,6 @@ def execute_query(
             query, db_args = _convert_to_old_pyformat(query, args, OperationalError)
             db_args.update(kwargs)
 
-        # Prefix each query transaction with experimental settings.
-        # These are necessary for nested datatypes to be available and other operations to work.
-        query = (
-            "set allow_experimental_lightweight_delete = 1;"
-            "set allow_experimental_object_type = 1;"
-            "set enable_http_compression= 1;"
-            f"{query}"
-        )
         with self._conn.cursor() as cursor:
             for query_line in query.split(";"):
                 if query_line := query_line.strip():
diff --git a/tests/load/clickhouse/test_clickhouse_configuration.py b/tests/load/clickhouse/test_clickhouse_configuration.py
index 61862170ed..5268c7bd22 100644
--- a/tests/load/clickhouse/test_clickhouse_configuration.py
+++ b/tests/load/clickhouse/test_clickhouse_configuration.py
@@ -1,11 +1,13 @@
-from typing import Any
+from typing import Any, Iterator
 
 import pytest
 
 import dlt
 from dlt.common.configuration.resolve import resolve_configuration
 from dlt.common.libs.sql_alchemy import make_url
+from dlt.common.storages import FileStorage
 from dlt.common.utils import digest128
+from dlt.destinations.impl.clickhouse.clickhouse import ClickHouseClient
 from dlt.destinations.impl.clickhouse.configuration import (
     ClickHouseCredentials,
     ClickHouseClientConfiguration,
@@ -15,6 +17,23 @@
     SnowflakeCredentials,
 )
 from tests.common.configuration.utils import environment
+from tests.load.utils import yield_client_with_storage
+from tests.utils import TEST_STORAGE_ROOT, delete_test_storage
+
+
+@pytest.fixture
+def file_storage() -> FileStorage:
+    return FileStorage(TEST_STORAGE_ROOT, file_type="b", makedirs=True)
+
+
+@pytest.fixture(autouse=True)
+def auto_delete_storage() -> None:
+    delete_test_storage()
+
+
+@pytest.fixture(scope="function")
+def client() -> Iterator[ClickHouseClient]:
+    yield from yield_client_with_storage("clickhouse")  # type: ignore
 
 
 def test_clickhouse_connection_string_with_all_params() -> None:
@@ -58,3 +77,20 @@ def test_clickhouse_gcp_hmac_getter_accessor(environment: Any) -> None:
     assert (
         dlt.config["destination.filesystem.credentials.gcp_secret_access_key"] == "ascvntp45uasdf"
     )
+
+
+def test_clickhouse_connection_settings(client: ClickHouseClient) -> None:
+    """Test experimental settings are set correctly for session."""
+    conn = client.sql_client.open_connection()
+    cursor1 = conn.cursor()
+    cursor2 = conn.cursor()
+
+    cursors = [cursor1, cursor2]
+
+    for cursor in cursors:
+        cursor.execute("SELECT name, value FROM system.settings")
+        res = cursor.fetchall()
+
+        assert ("allow_experimental_lightweight_delete", "1") in res
+        assert ("allow_experimental_object_type", "1") in res
+        assert ("enable_http_compression", "1") in res