diff --git a/sqlmesh/core/engine_adapter/clickhouse.py b/sqlmesh/core/engine_adapter/clickhouse.py index 92e3c33e6f..23dfb91ad2 100644 --- a/sqlmesh/core/engine_adapter/clickhouse.py +++ b/sqlmesh/core/engine_adapter/clickhouse.py @@ -61,6 +61,16 @@ 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 queries without the synthetic catalog unsupported by ClickHouse.""" + virtual_catalog = self._default_catalog or self._extra_config.get("virtual_catalog") + if virtual_catalog and isinstance(expression, (exp.Query, exp.Insert)): + expression = expression.copy() + for reference in expression.find_all(exp.Table, exp.Column): + if reference.text("catalog") == virtual_catalog: + reference.set("catalog", None) + return super()._to_sql(expression, quote=quote, **kwargs) + @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..6ed552075f 100644 --- a/tests/core/engine_adapter/test_clickhouse.py +++ b/tests/core/engine_adapter/test_clickhouse.py @@ -1596,6 +1596,75 @@ 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"', + ), + ], +) +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_execute_preserves_unconfigured_catalog_and_literals( + make_mocked_engine_adapter: t.Callable, +): + adapter = make_mocked_engine_adapter(ClickhouseEngineAdapter) + query = parse_one( + "SELECT other_catalog.mydb.source.id, '__ch_gw__.literal' FROM other_catalog.mydb.source", + dialect="clickhouse", + ) + + adapter.execute(query) + + assert to_sql_calls(adapter) == [ + 'SELECT "other_catalog"."mydb"."source"."id", \'__ch_gw__.literal\' ' + 'FROM "other_catalog"."mydb"."source"' + ] + + +def test_virtual_catalog_execute_uses_configured_catalog_fallback( + make_mocked_engine_adapter: t.Callable, +): + adapter = make_mocked_engine_adapter( + ClickhouseEngineAdapter, virtual_catalog="configured_catalog" + ) + query = parse_one( + "SELECT configured_catalog.mydb.source.id, other_catalog.otherdb.source.id, " + "'configured_catalog.literal' FROM configured_catalog.mydb.source " + "JOIN other_catalog.otherdb.source ON configured_catalog.mydb.source.id = " + "other_catalog.otherdb.source.id", + dialect="clickhouse", + ) + + adapter.execute(query) + + assert to_sql_calls(adapter) == [ + 'SELECT "mydb"."source"."id", "other_catalog"."otherdb"."source"."id", ' + '\'configured_catalog.literal\' FROM "mydb"."source" JOIN ' + '"other_catalog"."otherdb"."source" ON "mydb"."source"."id" = ' + '"other_catalog"."otherdb"."source"."id"' + ] + + def test_virtual_catalog_stripped_from_create_view_source( make_mocked_engine_adapter: t.Callable, ):