fix(clickhouse): strip virtual catalogs from executed queries and inserts

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"):
@@ -502,8 +517,14 @@ def _create_table_like(
         **kwargs: t.Any,
     ) -> None:
         """Create table with identical structure as source table"""
+        target_table_sql = self._strip_virtual_catalog(target_table_name).sql(
+            dialect=self.dialect, identify=True
+        )
+        source_table_sql = self._strip_virtual_catalog(source_table_name).sql(
+            dialect=self.dialect, identify=True
+        )
         self.execute(
-            f"CREATE TABLE {target_table_name}{self._on_cluster_sql()} AS {source_table_name}"
+            f"CREATE TABLE {target_table_sql}{self._on_cluster_sql()} AS {source_table_sql}"
         )
 
     def _get_partition_ids(
@@ -648,7 +669,7 @@ def _strip_virtual_catalog(self, name: "TableName") -> exp.Table:
         SQL is sent to the wire, since ClickHouse only supports a two-level
         ``[database].[table]`` naming scheme.
         """
-        table = exp.to_table(name)
+        table = exp.to_table(name, dialect=self.dialect)
         if self._default_catalog and table.catalog == self._default_catalog:
             table.set("catalog", None)
         return table
@@ -660,8 +681,12 @@ def _exchange_tables(
     ) -> None:
         from clickhouse_connect.driver.exceptions import DatabaseError  # type: ignore
 
-        old_table_sql = exp.to_table(old_table_name).sql(dialect=self.dialect, identify=True)
-        new_table_sql = exp.to_table(new_table_name).sql(dialect=self.dialect, identify=True)
+        old_table_sql = self._strip_virtual_catalog(old_table_name).sql(
+            dialect=self.dialect, identify=True
+        )
+        new_table_sql = self._strip_virtual_catalog(new_table_name).sql(
+            dialect=self.dialect, identify=True
+        )
 
         try:
             self.execute(
@@ -685,8 +710,12 @@ def _rename_table(
         old_table_name: TableName,
         new_table_name: TableName,
     ) -> None:
-        old_table_sql = exp.to_table(old_table_name).sql(dialect=self.dialect, identify=True)
-        new_table_sql = exp.to_table(new_table_name).sql(dialect=self.dialect, identify=True)
+        old_table_sql = self._strip_virtual_catalog(old_table_name).sql(
+            dialect=self.dialect, identify=True
+        )
+        new_table_sql = self._strip_virtual_catalog(new_table_name).sql(
+            dialect=self.dialect, identify=True
+        )
 
         self.execute(f"RENAME TABLE {old_table_sql} TO {new_table_sql}{self._on_cluster_sql()}")
 
@@ -974,7 +1003,7 @@ def _build_view_properties_exp(
     def _build_create_comment_table_exp(
         self, table: exp.Table, table_comment: str, table_kind: str, **kwargs: t.Any
     ) -> exp.Comment | str:
-        table_sql = table.sql(dialect=self.dialect, identify=True)
+        table_sql = self._strip_virtual_catalog(table).sql(dialect=self.dialect, identify=True)
 
         truncated_comment = self._truncate_table_comment(table_comment)
         comment_sql = exp.Literal.string(truncated_comment).sql(dialect=self.dialect)
@@ -989,7 +1018,7 @@ def _build_create_comment_column_exp(
         table_kind: str = "TABLE",
         **kwargs: t.Any,
     ) -> exp.Comment | str:
-        table_sql = table.sql(dialect=self.dialect, identify=True)
+        table_sql = self._strip_virtual_catalog(table).sql(dialect=self.dialect, identify=True)
         column_sql = exp.to_column(column_name).sql(dialect=self.dialect, identify=True)
 
         truncated_comment = self._truncate_table_comment(column_comment)

tests/core/engine_adapter/test_clickhouse.py

@@ -1066,7 +1066,7 @@ def test_insert_overwrite_by_condition_replace_partitioned(
     )
 
     assert to_sql_calls(adapter) == [
-        "CREATE TABLE __temp_target_abcd AS __temp_existing_records_abcd",
+        'CREATE TABLE "__temp_target_abcd" AS "__temp_existing_records_abcd"',
         'INSERT INTO "__temp_target_abcd" ("id", "ds") SELECT "id", "ds" FROM (SELECT * FROM "__temp_new_records_abcd") AS "_subquery"',
         'EXCHANGE TABLES "__temp_existing_records_abcd" AND "__temp_target_abcd"',
         'DROP TABLE IF EXISTS "__temp_target_abcd"',
@@ -1104,7 +1104,7 @@ def test_insert_overwrite_by_condition_replace(
     )
 
     to_sql_calls(adapter) == [
-        "CREATE TABLE __temp_target_abcd AS __temp_existing_records_abcd",
+        'CREATE TABLE "__temp_target_abcd" AS "__temp_existing_records_abcd"',
         'INSERT INTO "__temp_target_abcd" ("id", "ds") SELECT "id", "ds" FROM (SELECT * FROM "__temp_new_records_abcd") AS "_subquery"',
         'EXCHANGE TABLES "__temp_existing_records_abcd" AND "__temp_target_abcd"',
         'DROP TABLE IF EXISTS "__temp_target_abcd"',
@@ -1153,7 +1153,7 @@ def test_insert_overwrite_by_condition_where_partitioned(
     )
 
     to_sql_calls(adapter) == [
-        "CREATE TABLE __temp_target_abcd AS __temp_existing_records_abcd",
+        'CREATE TABLE "__temp_target_abcd" AS "__temp_existing_records_abcd"',
         """INSERT INTO "__temp_target_abcd" ("id", "ds") SELECT "id", "ds" FROM (SELECT * FROM "__temp_new_records_abcd") AS "_subquery" WHERE "ds" BETWEEN '2024-02-15' AND '2024-04-30'""",
         """CREATE TABLE IF NOT EXISTS "__temp_target_abcd" ENGINE=MergeTree ORDER BY () AS SELECT DISTINCT "partition_id" FROM (SELECT "_partition_id" AS "partition_id" FROM "__temp_existing_records_abcd" WHERE "ds" BETWEEN '2024-02-15' AND '2024-04-30' UNION DISTINCT SELECT "_partition_id" AS "partition_id" FROM "__temp_target_abcd") AS "_affected_partitions\"""",
         """INSERT INTO "__temp_target_abcd" SELECT "id", "ds" FROM "__temp_existing_records_abcd" WHERE NOT ("ds" BETWEEN '2024-02-15' AND '2024-04-30') AND "_partition_id" IN (SELECT "partition_id" FROM "__temp_target_abcd")""",
@@ -1204,12 +1204,12 @@ def test_insert_overwrite_by_condition_by_key(
     )
 
     to_sql_calls(adapter) == [
-        "CREATE TABLE __temp_target_abcd AS __temp_existing_records_abcd",
+        'CREATE TABLE "__temp_target_abcd" AS "__temp_existing_records_abcd"',
         'INSERT INTO "__temp_target_abcd" ("id", "ds") SELECT "id", "ds" FROM (SELECT DISTINCT ON ("id") * FROM "__temp_new_records_abcd") AS "_subquery"',
         'INSERT INTO "__temp_target_abcd" SELECT "id", "ds" FROM "__temp_existing_records_abcd" WHERE NOT ("id" IN (SELECT "id" FROM "__temp_target_abcd"))',
         'EXCHANGE TABLES "__temp_existing_records_abcd" AND "__temp_target_abcd"',
         'DROP TABLE IF EXISTS "__temp_target_abcd"',
-        "CREATE TABLE __temp_target_abcd AS __temp_existing_records_abcd",
+        'CREATE TABLE "__temp_target_abcd" AS "__temp_existing_records_abcd"',
         'INSERT INTO "__temp_target_abcd" ("id", "ds") SELECT "id", "ds" FROM (SELECT * FROM "__temp_new_records_abcd") AS "_subquery"',
         'INSERT INTO "__temp_target_abcd" SELECT "id", "ds" FROM "__temp_existing_records_abcd" WHERE NOT ("id" IN (SELECT "id" FROM "__temp_target_abcd"))',
         'EXCHANGE TABLES "__temp_existing_records_abcd" AND "__temp_target_abcd"',
@@ -1267,13 +1267,13 @@ def test_insert_overwrite_by_condition_by_key_partitioned(
     )
 
     to_sql_calls(adapter) == [
-        "CREATE TABLE __temp_target_abcd AS __temp_existing_records_abcd",
+        'CREATE TABLE "__temp_target_abcd" AS "__temp_existing_records_abcd"',
         'INSERT INTO "__temp_target_abcd" ("id", "ds") SELECT "id", "ds" FROM (SELECT DISTINCT ON ("id") * FROM "__temp_new_records_abcd") AS "_subquery"',
         'CREATE TABLE IF NOT EXISTS "__temp_target_abcd" ENGINE=MergeTree ORDER BY () AS SELECT DISTINCT "partition_id" FROM (SELECT "_partition_id" AS "partition_id" FROM "__temp_existing_records_abcd" WHERE "id" IN (SELECT "id" FROM "__temp_target_abcd") UNION DISTINCT SELECT "_partition_id" AS "partition_id" FROM "__temp_target_abcd") AS "_affected_partitions"',
         'INSERT INTO "__temp_target_abcd" SELECT "id", "ds" FROM "__temp_existing_records_abcd" WHERE NOT ("id" IN (SELECT "id" FROM "__temp_target_abcd")) AND "_partition_id" IN (SELECT "partition_id" FROM "__temp_target_abcd")',
         """ALTER TABLE "__temp_existing_records_abcd" REPLACE PARTITION ID '2' FROM "__temp_target_abcd", REPLACE PARTITION ID '1' FROM "__temp_target_abcd", REPLACE PARTITION ID '4' FROM "__temp_target_abcd", DROP PARTITION ID '3'""",
         'DROP TABLE IF EXISTS "__temp_target_abcd"',
-        "CREATE TABLE __temp_target_abcd AS __temp_existing_records_abcd",
+        'CREATE TABLE "__temp_target_abcd" AS "__temp_existing_records_abcd"',
         'INSERT INTO "__temp_target_abcd" ("id", "ds") SELECT "id", "ds" FROM (SELECT * FROM "__temp_new_records_abcd") AS "_subquery"',
         'CREATE TABLE IF NOT EXISTS "__temp_target_abcd" ENGINE=MergeTree ORDER BY () AS SELECT DISTINCT "partition_id" FROM (SELECT "_partition_id" AS "partition_id" FROM "__temp_existing_records_abcd" WHERE "id" IN (SELECT "id" FROM "__temp_target_abcd") UNION DISTINCT SELECT "_partition_id" AS "partition_id" FROM "__temp_target_abcd") AS "_affected_partitions"',
         'INSERT INTO "__temp_target_abcd" SELECT "id", "ds" FROM "__temp_existing_records_abcd" WHERE NOT ("id" IN (SELECT "id" FROM "__temp_target_abcd")) AND "_partition_id" IN (SELECT "partition_id" FROM "__temp_target_abcd")',
@@ -1316,7 +1316,7 @@ def test_insert_overwrite_by_condition_inc_by_partition(
     )
 
     to_sql_calls(adapter) == [
-        "CREATE TABLE __temp_target_abcd AS __temp_existing_records_abcd",
+        'CREATE TABLE "__temp_target_abcd" AS "__temp_existing_records_abcd"',
         'INSERT INTO "__temp_target_abcd" ("id", "ds") SELECT "id", "ds" FROM (SELECT * FROM "__temp_new_records_abcd") AS "_subquery"',
         """ALTER TABLE "__temp_existing_records_abcd" REPLACE PARTITION ID '1' FROM "__temp_target_abcd", REPLACE PARTITION ID '2' FROM "__temp_target_abcd", REPLACE PARTITION ID '4' FROM "__temp_target_abcd\"""",
         'DROP TABLE IF EXISTS "__temp_target_abcd"',
@@ -1596,6 +1596,124 @@ 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, '__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_insert_overwrite(
+    make_mocked_engine_adapter: t.Callable, mocker: MockerFixture
+):
+    adapter = make_mocked_engine_adapter(ClickhouseEngineAdapter)
+    adapter.inject_virtual_catalog("ch_gw")
+    mocker.patch(
+        "sqlmesh.core.engine_adapter.EngineAdapter._get_temp_table",
+        return_value=exp.to_table("__ch_gw__.mydb.__temp_target_abcd"),
+    )
+    mocker.patch("sqlmesh.core.engine_adapter.ClickhouseEngineAdapter.fetchone", return_value=None)
+
+    source_queries, columns_to_types = adapter._get_source_queries_and_columns_to_types(
+        parse_one("SELECT * FROM __ch_gw__.mydb.source"),
+        {"id": exp.DataType.build("Int8", dialect="clickhouse")},
+        "__ch_gw__.mydb.target",
+    )
+    adapter._insert_overwrite_by_condition(
+        "__ch_gw__.mydb.target", source_queries, columns_to_types
+    )
+
+    assert [call.args[0] for call in adapter.cursor.execute.call_args_list] == [
+        'CREATE TABLE "mydb"."__temp_target_abcd" AS "mydb"."target"',
+        'INSERT INTO "mydb"."__temp_target_abcd" ("id") SELECT "id" FROM '
+        '(SELECT * FROM "mydb"."source") AS "_subquery"',
+        'EXCHANGE TABLES "mydb"."target" AND "mydb"."__temp_target_abcd"',
+        'DROP TABLE IF EXISTS "mydb"."__temp_target_abcd"',
+    ]
+
+
+def test_virtual_catalog_stripped_from_rename_table(make_mocked_engine_adapter: t.Callable):
+    adapter = make_mocked_engine_adapter(ClickhouseEngineAdapter)
+    adapter.inject_virtual_catalog("ch_gw")
+
+    adapter.rename_table("__ch_gw__.mydb.old_table", "__ch_gw__.mydb.new_table")
+
+    assert [call.args[0] for call in adapter.cursor.execute.call_args_list] == [
+        'RENAME TABLE "mydb"."old_table" TO "mydb"."new_table"',
+    ]
+
+
+def test_virtual_catalog_stripped_from_comments(make_mocked_engine_adapter: t.Callable):
+    adapter = make_mocked_engine_adapter(ClickhouseEngineAdapter)
+    adapter.inject_virtual_catalog("ch_gw")
+
+    adapter._create_table_comment("__ch_gw__.mydb.target", "table comment")
+    adapter._create_column_comments("__ch_gw__.mydb.target", {"id": "column comment"})
+
+    assert [call.args[0] for call in adapter.cursor.execute.call_args_list] == [
+        'ALTER TABLE "mydb"."target" MODIFY COMMENT \'table comment\'',
+        'ALTER TABLE "mydb"."target" COMMENT COLUMN "id" \'column comment\'',
+    ]
+
+
+def test_three_part_names_unchanged_without_virtual_catalog(
+    make_mocked_engine_adapter: t.Callable,
+):
+    adapter = make_mocked_engine_adapter(ClickhouseEngineAdapter)
+
+    adapter.execute(parse_one("SELECT * FROM __ch_gw__.mydb.source", dialect="clickhouse"))
+
+    assert to_sql_calls(adapter) == ['SELECT * FROM "__ch_gw__"."mydb"."source"']
+
+
 def test_virtual_catalog_stripped_from_create_view_source(
     make_mocked_engine_adapter: t.Callable,
 ):