diff --git a/sqlmesh/core/engine_adapter/clickhouse.py b/sqlmesh/core/engine_adapter/clickhouse.py index 16ae5b5b51..70ea1807ec 100644 --- a/sqlmesh/core/engine_adapter/clickhouse.py +++ b/sqlmesh/core/engine_adapter/clickhouse.py @@ -61,6 +61,21 @@ def inject_virtual_catalog(self, gateway: str) -> None: configured = self._extra_config.get("virtual_catalog") self._default_catalog = f"__{gateway}__" if configured is None else configured + def _to_sql(self, expression: exp.Expr, quote: bool = True, **kwargs: t.Any) -> str: + """Render SQL without the virtual catalog, which ClickHouse does not support.""" + if self._default_catalog and any(self._virtual_catalog_references(expression)): + expression = expression.copy() + for reference in list(self._virtual_catalog_references(expression)): + reference.set("catalog", None) + return super()._to_sql(expression, quote=quote, **kwargs) + + def _virtual_catalog_references(self, expression: exp.Expr) -> t.Iterator[exp.Expr]: + return ( + reference + for reference in expression.find_all(exp.Table, exp.Column) + if reference.text("catalog") == self._default_catalog + ) + @property def engine_run_mode(self) -> EngineRunMode: if self._extra_config.get("cloud_mode"): diff --git a/tests/core/engine_adapter/test_clickhouse.py b/tests/core/engine_adapter/test_clickhouse.py index a3dfe0fdda..5af04304be 100644 --- a/tests/core/engine_adapter/test_clickhouse.py +++ b/tests/core/engine_adapter/test_clickhouse.py @@ -1596,6 +1596,65 @@ def test_virtual_catalog_stripped_in_alter_table(make_mocked_engine_adapter: t.C assert "ALTER TABLE" in sql_calls[0] +@pytest.mark.parametrize( + "query_sql, expected_sql", + [ + ( + 'INSERT INTO __ch_gw__.mydb.target ("id") ' + "SELECT __ch_gw__.mydb.source.id FROM __ch_gw__.mydb.source", + 'INSERT INTO "mydb"."target" ("id") SELECT "mydb"."source"."id" FROM "mydb"."source"', + ), + ( + "SELECT __ch_gw__.mydb.source.id FROM __ch_gw__.mydb.source", + 'SELECT "mydb"."source"."id" FROM "mydb"."source"', + ), + ( + "SELECT __ch_gw__.mydb.source.id, '__ch_gw__.literal' FROM __ch_gw__.mydb.source " + "JOIN other_catalog.otherdb.source ON __ch_gw__.mydb.source.id = " + "other_catalog.otherdb.source.id", + 'SELECT "mydb"."source"."id", \'__ch_gw__.literal\' FROM "mydb"."source" JOIN ' + '"other_catalog"."otherdb"."source" ON "mydb"."source"."id" = ' + '"other_catalog"."otherdb"."source"."id"', + ), + ], +) +def test_virtual_catalog_stripped_from_execute_queries( + make_mocked_engine_adapter: t.Callable, query_sql: str, expected_sql: str +): + adapter = make_mocked_engine_adapter(ClickhouseEngineAdapter) + adapter.inject_virtual_catalog("ch_gw") + query = parse_one(query_sql, dialect="clickhouse") + original_sql = query.sql(dialect="clickhouse") + + adapter.execute(query) + + assert query.sql(dialect="clickhouse") == original_sql + assert to_sql_calls(adapter) == [expected_sql] + + +def test_virtual_catalog_stripped_from_ctas_and_delete(make_mocked_engine_adapter: t.Callable): + adapter = make_mocked_engine_adapter(ClickhouseEngineAdapter) + adapter.inject_virtual_catalog("ch_gw") + + adapter.ctas( + "__ch_gw__.mydb.target", + parse_one("SELECT __ch_gw__.mydb.source.id FROM __ch_gw__.mydb.source"), + {"id": exp.DataType.build("Int32")}, + ) + adapter.delete_from( + "__ch_gw__.mydb.target", + "__ch_gw__.mydb.target.id IN (SELECT id FROM __ch_gw__.mydb.source)", + ) + + assert to_sql_calls(adapter) == [ + 'CREATE TABLE IF NOT EXISTS "mydb"."target" ENGINE=MergeTree ORDER BY () AS ' + 'SELECT CAST("id" AS Nullable(Int32)) AS "id" FROM ' + '(SELECT "mydb"."source"."id" FROM "mydb"."source") AS "_subquery"', + 'DELETE FROM "mydb"."target" WHERE "mydb"."target"."id" IN ' + '(SELECT "id" FROM "mydb"."source")', + ] + + def test_virtual_catalog_stripped_from_create_view_source( make_mocked_engine_adapter: t.Callable, ):