Feat!: add clickhouse engine adapter (#3019)

Co-authored-by: Erin Drummond <erin.dru@gmail.com>

Author: treysp

.circleci/continue_config.yml

@@ -315,6 +315,8 @@ workflows:
                 - mssql
                 - trino
                 - spark
+                - clickhouse
+                - clickhouse-cluster
       - engine_tests_cloud:
           name: cloud_engine_<< matrix.engine >>
           context:
@@ -328,6 +330,7 @@ workflows:
                 - databricks
                 - redshift
                 - bigquery
+                - clickhouse-cloud
           filters:
             branches:
               only:

.circleci/wait-for-db.sh

@@ -25,6 +25,10 @@ probe_port() {
     done
 }
 
+clickhouse_ready() {
+    probe_port 8123
+}
+
 postgres_ready() {
     probe_port 5432
 }

Makefile

@@ -10,7 +10,7 @@ install-doc:
 	pip3 install -r ./docs/requirements.txt
 
 install-engine-test:
-	pip3 install -e ".[dev,web,slack,mysql,postgres,databricks,redshift,bigquery,snowflake,trino,mssql]"
+	pip3 install -e ".[dev,web,slack,mysql,postgres,databricks,redshift,bigquery,snowflake,trino,mssql,clickhouse]"
 
 install-pre-commit:
 	pre-commit install
@@ -86,9 +86,9 @@ dev-publish: ui-build clean-build publish
 jupyter-example:
 	jupyter lab tests/slows/jupyter/example_outputs.ipynb
 
-engine-up: engine-mssql-up engine-mysql-up engine-postgres-up engine-spark-up engine-trino-up
+engine-up: engine-clickhouse-up engine-mssql-up engine-mysql-up engine-postgres-up engine-spark-up engine-trino-up
 
-engine-down: engine-mssql-down engine-mysql-down engine-postgres-down engine-spark-down engine-trino-down
+engine-down: engine-clickhouse-down engine-mssql-down engine-mysql-down engine-postgres-down engine-spark-down engine-trino-down
 
 fast-test:
 	pytest -n auto -m "fast and not cicdonly"
@@ -166,6 +166,12 @@ engine-%-down:
 # Docker Engines #
 ##################
 
+clickhouse-test: engine-clickhouse-up
+	pytest -n auto -x -m "clickhouse" --junitxml=test-results/junit-clickhouse.xml
+
+clickhouse-cluster-test: engine-clickhouse-up
+	pytest -n auto -x -m "clickhouse_cluster" --junitxml=test-results/junit-clickhouse-cluster.xml
+
 duckdb-test: engine-duckdb-install
 	pytest -n auto -x -m "duckdb" --junitxml=test-results/junit-duckdb.xml
 
@@ -200,3 +206,6 @@ databricks-test: guard-DATABRICKS_CATALOG guard-DATABRICKS_SERVER_HOSTNAME guard
 
 redshift-test: guard-REDSHIFT_HOST guard-REDSHIFT_USER guard-REDSHIFT_PASSWORD guard-REDSHIFT_DATABASE engine-redshift-install
 	pytest -n auto -x -m "redshift" --junitxml=test-results/junit-redshift.xml
+
+clickhouse-cloud-test: guard-CLICKHOUSE_CLOUD_HOST guard-CLICKHOUSE_CLOUD_USERNAME guard-CLICKHOUSE_CLOUD_PASSWORD engine-clickhouse-install
+	pytest -n auto -x -m "clickhouse_cloud" --junitxml=test-results/junit-clickhouse-cloud.xml

examples/sushi/models/raw_marketing.py

@@ -59,7 +59,15 @@ def execute(
             "updated_at": [exec_time] * num_customers,
         }
     )
-    df = df_new.merge(df_existing, on="customer_id", how="left", suffixes=(None, "_old"))
+
+    # clickhouse returns a dataframe with no columns if the query is empty, so we can't merge
+    if not df_existing.empty:
+        df = df_new.merge(df_existing, on="customer_id", how="left", suffixes=(None, "_old"))
+    else:
+        df = df_new
+        df["status_old"] = pd.NA
+        df["updated_at_old"] = pd.NA
+
     df["updated_at"] = pd.to_datetime(
         np.where(  # type: ignore
             df["status_old"] != df["status"], execution_time, df["updated_at_old"]

pytest.ini

@@ -20,6 +20,9 @@ markers =
     # Engine Adapters
     engine: test all engine adapters
     bigquery: test for BigQuery
+    clickhouse: test for Clickhouse (standalone mode)
+    clickhouse_cluster: test for Clickhouse (cluster mode)
+    clickhouse_cloud: test for Clickhouse (cloud mode)
     databricks: test for Databricks
     duckdb: test for DuckDB
     motherduck: test for MotherDuck

setup.py

@@ -47,13 +47,14 @@
         "requests",
         "rich[jupyter]",
         "ruamel.yaml",
-        "sqlglot[rs]~=25.17.0",
+        "sqlglot[rs]~=25.18.0",
     ],
     extras_require={
         "bigquery": [
             "google-cloud-bigquery[pandas]",
             "google-cloud-bigquery-storage",
         ],
+        "clickhouse": ["clickhouse-connect"],
         "databricks": ["databricks-sql-connector"],
         "dev": [
             f"apache-airflow=={os.environ.get('AIRFLOW_VERSION', '2.9.1')}",

sqlmesh/core/config/connection.py

@@ -39,6 +39,13 @@
 logger = logging.getLogger(__name__)
 
 RECOMMENDED_STATE_SYNC_ENGINES = {"postgres", "gcp_postgres", "mysql", "duckdb"}
+FORBIDDEN_STATE_SYNC_ENGINES = {
+    # Do not support row-level operations
+    "spark",
+    "trino",
+    # Nullable types are problematic
+    "clickhouse",
+}
 
 
 class ConnectionConfig(abc.ABC, BaseConfig):
@@ -84,9 +91,14 @@ def _cursor_init(self) -> t.Optional[t.Callable[[t.Any], None]]:
 
     @property
     def is_recommended_for_state_sync(self) -> bool:
-        """Whether this connection is recommended for being used as a state sync for production state syncs"""
+        """Whether this engine is recommended for being used as a state sync for production state syncs"""
         return self.type_ in RECOMMENDED_STATE_SYNC_ENGINES
 
+    @property
+    def is_forbidden_for_state_sync(self) -> bool:
+        """Whether this engine is forbidden from being used as a state sync"""
+        return self.type_ in FORBIDDEN_STATE_SYNC_ENGINES
+
     @property
     def _connection_factory_with_kwargs(self) -> t.Callable[[], t.Any]:
         """A function that is called to return a connection object for the given Engine Adapter"""
@@ -1364,6 +1376,73 @@ def _static_connection_kwargs(self) -> t.Dict[str, t.Any]:
         }
 
 
+class ClickhouseConnectionConfig(ConnectionConfig):
+    """
+    Clickhouse Connection Configuration.
+
+    Property reference: https://clickhouse.com/docs/en/integrations/python#client-initialization
+    """
+
+    host: str
+    username: str
+    password: t.Optional[str] = None
+    port: t.Optional[int] = None
+    cluster: t.Optional[str] = None
+    connect_timeout: int = 10
+    send_receive_timeout: int = 300
+    verify: bool = True
+    query_limit: int = 0
+    use_compression: bool = True
+    compression_method: t.Optional[str] = None
+
+    concurrent_tasks: int = 1
+    register_comments: bool = True
+    pre_ping: bool = False
+
+    type_: Literal["clickhouse"] = Field(alias="type", default="clickhouse")
+
+    @property
+    def _connection_kwargs_keys(self) -> t.Set[str]:
+        kwargs = {
+            "host",
+            "username",
+            "port",
+            "password",
+            "connect_timeout",
+            "send_receive_timeout",
+            "verify",
+            "query_limit",
+        }
+        return kwargs
+
+    @property
+    def _engine_adapter(self) -> t.Type[EngineAdapter]:
+        return engine_adapter.ClickhouseEngineAdapter
+
+    @property
+    def _connection_factory(self) -> t.Callable:
+        from clickhouse_connect.dbapi import connect  # type: ignore
+
+        return connect
+
+    @property
+    def _extra_engine_config(self) -> t.Dict[str, t.Any]:
+        return {"cluster": self.cluster}
+
+    @property
+    def _static_connection_kwargs(self) -> t.Dict[str, t.Any]:
+        from sqlmesh import __version__
+
+        # False = no compression
+        # True = Clickhouse default compression method
+        # string = specific compression method
+        compress: bool | str = self.use_compression
+        if compress and self.compression_method:
+            compress = self.compression_method
+
+        return {"compress": compress, "client_name": f"SQLMesh/{__version__}"}
+
+
 CONNECTION_CONFIG_TO_TYPE = {
     # Map all subclasses of ConnectionConfig to the value of their `type_` field.
     tpe.all_field_infos()["type_"].default: tpe

sqlmesh/core/config/scheduler.py

@@ -83,7 +83,7 @@ def create_state_sync(self, context: GenericContext) -> StateSync:
             context.config.get_state_connection(context.gateway) or context._connection_config
         )
         engine_adapter = state_connection.create_engine_adapter()
-        if not engine_adapter.SUPPORTS_ROW_LEVEL_OP:
+        if state_connection.is_forbidden_for_state_sync:
             raise ConfigError(
                 f"The {engine_adapter.DIALECT.upper()} engine cannot be used to store SQLMesh state - please specify a different `state_connection` engine."
                 + " See https://sqlmesh.readthedocs.io/en/stable/reference/configuration/#gateways for more information."

sqlmesh/core/dialect.py

@@ -374,6 +374,9 @@ def _parse_props(self: Parser) -> t.Optional[exp.Expression]:
     name = key.name.lower()
     if name == "when_matched":
         value: t.Optional[exp.Expression] = self._parse_when_matched()[0]
+    elif name == "time_data_type":
+        # TODO: if we make *_data_type a convention to parse things into exp.DataType, we could make this more generic
+        value = self._parse_types(schema=True)
     elif self._match(TokenType.L_PAREN):
         value = self.expression(exp.Tuple, expressions=self._parse_csv(self._parse_equality))
         self._match_r_paren()

sqlmesh/core/engine_adapter/__init__.py

@@ -7,6 +7,7 @@
     EngineAdapterWithIndexSupport,
 )
 from sqlmesh.core.engine_adapter.bigquery import BigQueryEngineAdapter
+from sqlmesh.core.engine_adapter.clickhouse import ClickhouseEngineAdapter
 from sqlmesh.core.engine_adapter.databricks import DatabricksEngineAdapter
 from sqlmesh.core.engine_adapter.duckdb import DuckDBEngineAdapter
 from sqlmesh.core.engine_adapter.mssql import MSSQLEngineAdapter
@@ -21,6 +22,7 @@
     "hive": SparkEngineAdapter,
     "spark": SparkEngineAdapter,
     "bigquery": BigQueryEngineAdapter,
+    "clickhouse": ClickhouseEngineAdapter,
     "duckdb": DuckDBEngineAdapter,
     "snowflake": SnowflakeEngineAdapter,
     "databricks": DatabricksEngineAdapter,