diff --git a/sqlmesh/core/engine_adapter/clickhouse.py b/sqlmesh/core/engine_adapter/clickhouse.py index e385452bf3..b495abf78a 100644 --- a/sqlmesh/core/engine_adapter/clickhouse.py +++ b/sqlmesh/core/engine_adapter/clickhouse.py @@ -446,7 +446,7 @@ def _build_alter_partition_exp( ) -> exp.Alter: alter_expr = exp.Alter(this=target_table, kind="TABLE") - for partition in partitions_to_replace: + for partition in sorted(partitions_to_replace): alter_expr.append( "actions", exp.ReplacePartition( @@ -457,7 +457,7 @@ def _build_alter_partition_exp( ), ) - for partition in partitions_to_drop: + for partition in sorted(partitions_to_drop): alter_expr.append( "actions", exp.DropPartition( diff --git a/tests/core/engine_adapter/test_clickhouse.py b/tests/core/engine_adapter/test_clickhouse.py index e0355f456d..91f287722e 100644 --- a/tests/core/engine_adapter/test_clickhouse.py +++ b/tests/core/engine_adapter/test_clickhouse.py @@ -1118,7 +1118,7 @@ def test_insert_overwrite_by_condition_replace( columns_to_types, ) - to_sql_calls(adapter) == [ + assert to_sql_calls(adapter) == [ '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"', @@ -1133,7 +1133,10 @@ def test_insert_overwrite_by_condition_where_partitioned( temp_table_mock = mocker.patch("sqlmesh.core.engine_adapter.EngineAdapter._get_temp_table") table_name = "target" - temp_table_mock.return_value = make_temp_table_name(table_name, "abcd") + temp_table_mock.side_effect = lambda table, *args, **kwargs: make_temp_table_name( + "affected_partitions" if exp.to_table(table).name == "_affected_partitions" else table_name, + "abcd", + ) fetchone_mock = mocker.patch("sqlmesh.core.engine_adapter.ClickhouseEngineAdapter.fetchone") fetchone_mock.return_value = "dateTrunc('WEEK', ds)" @@ -1141,7 +1144,7 @@ def test_insert_overwrite_by_condition_where_partitioned( fetchall_mock = mocker.patch("sqlmesh.core.engine_adapter.ClickhouseEngineAdapter.fetchall") fetchall_mock.side_effect = [ [("1",), ("2",), ("3",), ("4",)], - ["1", "2", "4"], + [("1",), ("2",), ("4",)], ] insert_table_name = make_temp_table_name("new_records", "abcd") @@ -1167,11 +1170,12 @@ def test_insert_overwrite_by_condition_where_partitioned( ), ) - to_sql_calls(adapter) == [ + assert to_sql_calls(adapter) == [ '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")""", + '''CREATE TABLE IF NOT EXISTS "__temp_affected_partitions_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_affected_partitions_abcd")""", + 'DROP TABLE IF EXISTS "__temp_affected_partitions_abcd"', """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 PARTITION ID '3'""", 'DROP TABLE IF EXISTS "__temp_target_abcd"', ] @@ -1218,7 +1222,7 @@ def test_insert_overwrite_by_condition_by_key( dynamic_key_unique=False, ) - to_sql_calls(adapter) == [ + assert to_sql_calls(adapter) == [ '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"))', @@ -1226,7 +1230,7 @@ def test_insert_overwrite_by_condition_by_key( 'DROP TABLE IF EXISTS "__temp_target_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"))', + 'INSERT INTO "__temp_target_abcd" SELECT "id", "ds" FROM "__temp_existing_records_abcd" WHERE NOT ("id" IN (SELECT DISTINCT "id" FROM "__temp_target_abcd"))', 'EXCHANGE TABLES "__temp_existing_records_abcd" AND "__temp_target_abcd"', 'DROP TABLE IF EXISTS "__temp_target_abcd"', ] @@ -1239,7 +1243,10 @@ def test_insert_overwrite_by_condition_by_key_partitioned( temp_table_mock = mocker.patch("sqlmesh.core.engine_adapter.EngineAdapter._get_temp_table") table_name = "target" - temp_table_mock.return_value = make_temp_table_name(table_name, "abcd") + temp_table_mock.side_effect = lambda table, *args, **kwargs: make_temp_table_name( + "affected_partitions" if exp.to_table(table).name == "_affected_partitions" else table_name, + "abcd", + ) fetchone_mock = mocker.patch("sqlmesh.core.engine_adapter.ClickhouseEngineAdapter.fetchone") fetchone_mock.side_effect = ["dateTrunc('WEEK', ds)", "dateTrunc('WEEK', ds)"] @@ -1247,9 +1254,9 @@ def test_insert_overwrite_by_condition_by_key_partitioned( fetchall_mock = mocker.patch("sqlmesh.core.engine_adapter.ClickhouseEngineAdapter.fetchall") fetchall_mock.side_effect = [ [("1",), ("2",), ("3",), ("4",)], - ["1", "2", "4"], + [("1",), ("2",), ("4",)], [("1",), ("2",), ("3",), ("4",)], - ["1", "2", "4"], + [("1",), ("2",), ("4",)], ] insert_table_name = make_temp_table_name("new_records", "abcd") @@ -1281,18 +1288,20 @@ def test_insert_overwrite_by_condition_by_key_partitioned( dynamic_key_unique=False, ) - to_sql_calls(adapter) == [ + assert to_sql_calls(adapter) == [ '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'""", + 'CREATE TABLE IF NOT EXISTS "__temp_affected_partitions_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_affected_partitions_abcd")', + 'DROP TABLE IF EXISTS "__temp_affected_partitions_abcd"', + """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 PARTITION ID '3'""", 'DROP TABLE IF EXISTS "__temp_target_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")', - """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'""", + 'CREATE TABLE IF NOT EXISTS "__temp_affected_partitions_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 DISTINCT "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 DISTINCT "id" FROM "__temp_target_abcd")) AND "_partition_id" IN (SELECT "partition_id" FROM "__temp_affected_partitions_abcd")', + 'DROP TABLE IF EXISTS "__temp_affected_partitions_abcd"', + """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 PARTITION ID '3'""", 'DROP TABLE IF EXISTS "__temp_target_abcd"', ] @@ -1330,7 +1339,7 @@ def test_insert_overwrite_by_condition_inc_by_partition( keep_existing_partition_rows=False, ) - to_sql_calls(adapter) == [ + assert to_sql_calls(adapter) == [ '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\"""",