diff --git a/sqlmesh/core/engine_adapter/clickhouse.py b/sqlmesh/core/engine_adapter/clickhouse.py index 16ae5b5b51..e64cac9f0e 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"): @@ -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) diff --git a/tests/core/engine_adapter/test_clickhouse.py b/tests/core/engine_adapter/test_clickhouse.py index a3dfe0fdda..38079302d9 100644 --- a/tests/core/engine_adapter/test_clickhouse.py +++ b/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, ):