Skip to content

Commit

Permalink
Add test for clickhouse config settings #1055
Browse files Browse the repository at this point in the history
Signed-off-by: Marcel Coetzee <[email protected]>
  • Loading branch information
Pipboyguy committed Apr 16, 2024
1 parent 974d150 commit 027d52a
Show file tree
Hide file tree
Showing 3 changed files with 39 additions and 11 deletions.
1 change: 1 addition & 0 deletions dlt/destinations/impl/clickhouse/clickhouse.py
Original file line number Diff line number Diff line change
Expand Up @@ -254,6 +254,7 @@ def __init__(
settings={
"allow_experimental_lightweight_delete": 1,
"allow_experimental_object_type": 1,
"enable_http_compression": 1,
},
compression=compression,
)
Expand Down
11 changes: 1 addition & 10 deletions dlt/destinations/impl/clickhouse/sql_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -128,22 +127,14 @@ 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()

if args:
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():
Expand Down
38 changes: 37 additions & 1 deletion tests/load/clickhouse/test_clickhouse_configuration.py
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -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:
Expand Down Expand Up @@ -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

0 comments on commit 027d52a

Please sign in to comment.