Support secure connection to Clickhouse

docsource/usage/clickhouse/settings.md

@@ -10,4 +10,7 @@ In order to use SQL Mock with Clickhouse, you need to provide the following envi
 * `SQL_MOCK_CLICKHOUSE_PASSWORD`: Password of your user
 * `SQL_MOCK_CLICKHOUSE_PORT`: Port of your Clickhouse instance
 
+Additionally, there are optional environment variables:
+* `SQL_MOCK_CLICKHOUSE_USE_SECURE_CONNECTION`: Whether to use a secure connection or not (default False)
+
 Having those environment variables enables SQL Mock to connect to your Clickhouse instance.

pyproject.toml

@@ -43,7 +43,8 @@ pydantic-settings = "^2.0.3"
 sqlglot = "^20.5.0"
 
 # Clickhouse specific
-clickhouse-driver = {extras = ["numpy"], version = "^0.2.6", optional = true}
+clickhouse-connect = {version = "^0.7.0", optional = true}
+numpy = {version = "^1.26.3", optional = true}
 
 # Google Bigquery specific
 google-cloud-bigquery = {version ="^3.11.4", optional = true}
@@ -57,7 +58,7 @@ snowflake-connector-python = "^3.6.0"
 
 [tool.poetry.extras]
 bigquery = ["google-cloud-bigquery"]
-clickhouse = ["clickhouse-driver"]
+clickhouse = ["clickhouse-connect", "numpy"]
 redshift = ["redshift-connector", "boto3"]
 snowflake = ["snowflake-connector-python"]
 

src/sql_mock/clickhouse/settings.py

@@ -7,3 +7,4 @@ class ClickHouseSettings(BaseSettings):
     user: str
     password: str
     port: str
+    use_secure_connection: bool = False

src/sql_mock/clickhouse/table_mocks.py

@@ -1,4 +1,4 @@
-from clickhouse_driver import Client
+import clickhouse_connect
 
 from sql_mock.clickhouse.settings import ClickHouseSettings
 from sql_mock.table_mocks import BaseMockTable
@@ -16,8 +16,12 @@ def __init__(
         super().__init__(*args, **kwargs)
 
     def _get_results(self, query: str) -> list[dict]:
-        with Client(
-            host=self.settings.host, user=self.settings.user, password=self.settings.password, port=self.settings.port
+        with clickhouse_connect.get_client(
+            host=self.settings.host,
+            secure=self.settings.use_secure_connection,
+            username=self.settings.user,
+            password=self.settings.password,
+            port=self.settings.port,
         ) as client:
-            res = client.query_dataframe(query)
+            res = client.query_df(query)
         return res.to_dict("records")

src/sql_mock/table_mocks.py

@@ -168,7 +168,9 @@ def _generate_query(
         input_data_ctes = self._generate_input_data_cte_snippet()
 
         # Parse the query with sqlglot to to standardize it (e.g. removes semi-colons)
-        result_query = sqlglot.parse_one(self._sql_mock_data.rendered_query, dialect=self._sql_dialect).sql()
+        result_query = sqlglot.parse_one(self._sql_mock_data.rendered_query, dialect=self._sql_dialect).sql(
+            dialect=self._sql_dialect
+        )
 
         if cte_to_select is not None:
             result_query = select_from_cte(result_query, cte_to_select, sql_dialect=self._sql_dialect)

tests/sql_mock/clickhouse/test_table_mocks.py

@@ -55,16 +55,16 @@ def test_get_results(mocker):
     """
     Test the _get_results method.
     """
-    mock_client = mocker.patch("sql_mock.clickhouse.table_mocks.Client")
+    mock_client = mocker.patch("sql_mock.clickhouse.table_mocks.clickhouse_connect.get_client")
     mock_query_result = [{"column1": "value1", "column2": 42}]
     query = "SELECT 1, 2"
 
     mock_dataframe = mocker.MagicMock()
     mock_dataframe.to_dict.return_value = mock_query_result
-    mock_client.return_value.__enter__.return_value.query_dataframe.return_value = mock_dataframe
+    mock_client.return_value.__enter__.return_value.query_df.return_value = mock_dataframe
 
     instance = ClickHouseTableMock()
     result = instance._get_results(query=query)
 
     assert result == mock_query_result
-    mock_client.return_value.__enter__.return_value.query_dataframe.assert_called_once_with(query)
+    mock_client.return_value.__enter__.return_value.query_df.assert_called_once_with(query)