fix(clickhouse): add clickhouse connect driver (#23185)

2026-04-20 00:24:38 +00:00 · 2023-02-24 14:04:12 +02:00
parent f0f27a486d
commit d0c54cddb0
4 changed files with 456 additions and 29 deletions
--- a/tests/unit_tests/db_engine_specs/test_clickhouse.py
+++ b/tests/unit_tests/db_engine_specs/test_clickhouse.py
@@ -16,12 +16,26 @@
 # under the License.

 from datetime import datetime
-from typing import Optional
+from typing import Any, Dict, Optional, Type
 from unittest.mock import Mock

 import pytest
+from sqlalchemy.types import (
+    Boolean,
+    Date,
+    DateTime,
+    DECIMAL,
+    Float,
+    Integer,
+    String,
+    TypeEngine,
+)

-from tests.unit_tests.db_engine_specs.utils import assert_convert_dttm
+from superset.utils.core import GenericDataType
+from tests.unit_tests.db_engine_specs.utils import (
+    assert_column_spec,
+    assert_convert_dttm,
+)
 from tests.unit_tests.fixtures.common import dttm


@@ -53,3 +67,147 @@ def test_execute_connection_error() -> None:
    )
    with pytest.raises(SupersetDBAPIDatabaseError) as ex:
        ClickHouseEngineSpec.execute(cursor, "SELECT col1 from table1")
+
+
+@pytest.mark.parametrize(
+    "target_type,expected_result",
+    [
+        ("Date", "toDate('2019-01-02')"),
+        ("DateTime", "toDateTime('2019-01-02 03:04:05')"),
+        ("UnknownType", None),
+    ],
+)
+def test_connect_convert_dttm(
+    target_type: str, expected_result: Optional[str], dttm: datetime
+) -> None:
+    from superset.db_engine_specs.clickhouse import ClickHouseEngineSpec as spec
+
+    assert_convert_dttm(spec, target_type, expected_result, dttm)
+
+
+@pytest.mark.parametrize(
+    "native_type,sqla_type,attrs,generic_type,is_dttm",
+    [
+        ("String", String, None, GenericDataType.STRING, False),
+        ("LowCardinality(String)", String, None, GenericDataType.STRING, False),
+        ("Nullable(String)", String, None, GenericDataType.STRING, False),
+        (
+            "LowCardinality(Nullable(String))",
+            String,
+            None,
+            GenericDataType.STRING,
+            False,
+        ),
+        ("Array(UInt8)", String, None, GenericDataType.STRING, False),
+        ("Enum('hello', 'world')", String, None, GenericDataType.STRING, False),
+        ("Enum('UInt32', 'Bool')", String, None, GenericDataType.STRING, False),
+        (
+            "LowCardinality(Enum('hello', 'world'))",
+            String,
+            None,
+            GenericDataType.STRING,
+            False,
+        ),
+        (
+            "Nullable(Enum('hello', 'world'))",
+            String,
+            None,
+            GenericDataType.STRING,
+            False,
+        ),
+        (
+            "LowCardinality(Nullable(Enum('hello', 'world')))",
+            String,
+            None,
+            GenericDataType.STRING,
+            False,
+        ),
+        ("FixedString(16)", String, None, GenericDataType.STRING, False),
+        ("Nullable(FixedString(16))", String, None, GenericDataType.STRING, False),
+        (
+            "LowCardinality(Nullable(FixedString(16)))",
+            String,
+            None,
+            GenericDataType.STRING,
+            False,
+        ),
+        ("UUID", String, None, GenericDataType.STRING, False),
+        ("Int8", Integer, None, GenericDataType.NUMERIC, False),
+        ("Int16", Integer, None, GenericDataType.NUMERIC, False),
+        ("Int32", Integer, None, GenericDataType.NUMERIC, False),
+        ("Int64", Integer, None, GenericDataType.NUMERIC, False),
+        ("Int128", Integer, None, GenericDataType.NUMERIC, False),
+        ("Int256", Integer, None, GenericDataType.NUMERIC, False),
+        ("Nullable(Int256)", Integer, None, GenericDataType.NUMERIC, False),
+        (
+            "LowCardinality(Nullable(Int256))",
+            Integer,
+            None,
+            GenericDataType.NUMERIC,
+            False,
+        ),
+        ("UInt8", Integer, None, GenericDataType.NUMERIC, False),
+        ("UInt16", Integer, None, GenericDataType.NUMERIC, False),
+        ("UInt32", Integer, None, GenericDataType.NUMERIC, False),
+        ("UInt64", Integer, None, GenericDataType.NUMERIC, False),
+        ("UInt128", Integer, None, GenericDataType.NUMERIC, False),
+        ("UInt256", Integer, None, GenericDataType.NUMERIC, False),
+        ("Nullable(UInt256)", Integer, None, GenericDataType.NUMERIC, False),
+        (
+            "LowCardinality(Nullable(UInt256))",
+            Integer,
+            None,
+            GenericDataType.NUMERIC,
+            False,
+        ),
+        ("Float32", Float, None, GenericDataType.NUMERIC, False),
+        ("Float64", Float, None, GenericDataType.NUMERIC, False),
+        ("Decimal(1, 2)", DECIMAL, None, GenericDataType.NUMERIC, False),
+        ("Decimal32(2)", DECIMAL, None, GenericDataType.NUMERIC, False),
+        ("Decimal64(2)", DECIMAL, None, GenericDataType.NUMERIC, False),
+        ("Decimal128(2)", DECIMAL, None, GenericDataType.NUMERIC, False),
+        ("Decimal256(2)", DECIMAL, None, GenericDataType.NUMERIC, False),
+        ("Bool", Boolean, None, GenericDataType.BOOLEAN, False),
+        ("Nullable(Bool)", Boolean, None, GenericDataType.BOOLEAN, False),
+        ("Date", Date, None, GenericDataType.TEMPORAL, True),
+        ("Nullable(Date)", Date, None, GenericDataType.TEMPORAL, True),
+        ("LowCardinality(Nullable(Date))", Date, None, GenericDataType.TEMPORAL, True),
+        ("Date32", Date, None, GenericDataType.TEMPORAL, True),
+        ("Datetime", DateTime, None, GenericDataType.TEMPORAL, True),
+        ("Nullable(Datetime)", DateTime, None, GenericDataType.TEMPORAL, True),
+        (
+            "LowCardinality(Nullable(Datetime))",
+            DateTime,
+            None,
+            GenericDataType.TEMPORAL,
+            True,
+        ),
+        ("Datetime('UTC')", DateTime, None, GenericDataType.TEMPORAL, True),
+        ("Datetime64(3)", DateTime, None, GenericDataType.TEMPORAL, True),
+        ("Datetime64(3, 'UTC')", DateTime, None, GenericDataType.TEMPORAL, True),
+    ],
+)
+def test_connect_get_column_spec(
+    native_type: str,
+    sqla_type: Type[TypeEngine],
+    attrs: Optional[Dict[str, Any]],
+    generic_type: GenericDataType,
+    is_dttm: bool,
+) -> None:
+    from superset.db_engine_specs.clickhouse import ClickHouseConnectEngineSpec as spec
+
+    assert_column_spec(spec, native_type, sqla_type, attrs, generic_type, is_dttm)
+
+
+@pytest.mark.parametrize(
+    "column_name,expected_result",
+    [
+        ("time", "time_07cc69"),
+        ("count", "count_e2942a"),
+    ],
+)
+def test_connect_make_label_compatible(column_name: str, expected_result: str) -> None:
+    from superset.db_engine_specs.clickhouse import ClickHouseConnectEngineSpec as spec
+
+    label = spec.make_label_compatible(column_name)
+    assert label == expected_result
--- a/tests/unit_tests/db_engine_specs/test_mssql.py
+++ b/tests/unit_tests/db_engine_specs/test_mssql.py
@@ -20,7 +20,7 @@ from textwrap import dedent
 from typing import Any, Dict, Optional, Type

 import pytest
-from sqlalchemy import column, table, types
+from sqlalchemy import column, table
 from sqlalchemy.dialects import mssql
 from sqlalchemy.dialects.mssql import DATE, NTEXT, NVARCHAR, TEXT, VARCHAR
 from sqlalchemy.sql import select
@@ -50,7 +50,7 @@ from tests.unit_tests.fixtures.common import dttm
 )
 def test_get_column_spec(
    native_type: str,
-    sqla_type: Type[types.TypeEngine],
+    sqla_type: Type[TypeEngine],
    attrs: Optional[Dict[str, Any]],
    generic_type: GenericDataType,
    is_dttm: bool,