From e3f3b7a606b61ee235cab7db2467109491ec46af Mon Sep 17 00:00:00 2001 From: Michael Day Date: Fri, 16 Oct 2026 16:55:41 -0400 Subject: [PATCH] test: assert ClickHouse overwrite SQL Signed-off-by: Michael Day --- tests/core/engine_adapter/test_clickhouse.py | 106 ++++++++++++++----- 1 file changed, 81 insertions(+), 25 deletions(-) diff --git a/tests/core/engine_adapter/test_clickhouse.py b/tests/core/engine_adapter/test_clickhouse.py index e0355f456d..4e80393604 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,15 +1170,29 @@ def test_insert_overwrite_by_condition_where_partitioned( ), ) - to_sql_calls(adapter) == [ + sql_calls = to_sql_calls(adapter) + + assert len(sql_calls) == 7 + assert sql_calls[:5] == [ '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")""", - """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 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_prefix = 'ALTER TABLE "__temp_existing_records_abcd" ' + assert sql_calls[5].startswith(alter_prefix) + assert sorted(sql_calls[5].removeprefix(alter_prefix).split(", ")) == sorted( + [ + "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'", + ] + ) + assert sql_calls[6] == 'DROP TABLE IF EXISTS "__temp_target_abcd"' + def test_insert_overwrite_by_condition_by_key( make_mocked_engine_adapter: t.Callable, mocker: MockerFixture, make_temp_table_name: t.Callable @@ -1218,7 +1235,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 +1243,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 +1256,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 +1267,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,20 +1301,45 @@ def test_insert_overwrite_by_condition_by_key_partitioned( dynamic_key_unique=False, ) - to_sql_calls(adapter) == [ + sql_calls = to_sql_calls(adapter) + + assert len(sql_calls) == 14 + assert sql_calls[:4] == [ '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 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")', + ] + + assert sql_calls[4] == 'DROP TABLE IF EXISTS "__temp_affected_partitions_abcd"' + alter_prefix = 'ALTER TABLE "__temp_existing_records_abcd" ' + assert sql_calls[5].startswith(alter_prefix) + assert sorted(sql_calls[5].removeprefix(alter_prefix).split(", ")) == sorted( + [ + "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'", + ] + ) + assert sql_calls[6] == 'DROP TABLE IF EXISTS "__temp_target_abcd"' + assert sql_calls[7:11] == [ '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'""", - 'DROP TABLE IF EXISTS "__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 "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")', ] + assert sql_calls[11] == 'DROP TABLE IF EXISTS "__temp_affected_partitions_abcd"' + assert sql_calls[12].startswith(alter_prefix) + assert sorted(sql_calls[12].removeprefix(alter_prefix).split(", ")) == sorted( + [ + "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'", + ] + ) + assert sql_calls[13] == 'DROP TABLE IF EXISTS "__temp_target_abcd"' def test_insert_overwrite_by_condition_inc_by_partition( @@ -1330,12 +1375,23 @@ def test_insert_overwrite_by_condition_inc_by_partition( keep_existing_partition_rows=False, ) - to_sql_calls(adapter) == [ + sql_calls = to_sql_calls(adapter) + + assert len(sql_calls) == 4 + assert sql_calls[:2] == [ '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"', ] + alter_prefix = 'ALTER TABLE "__temp_existing_records_abcd" ' + assert sql_calls[2].startswith(alter_prefix) + assert sorted(sql_calls[2].removeprefix(alter_prefix).split(", ")) == sorted( + [ + "REPLACE PARTITION ID '1' FROM \"__temp_target_abcd\"", + "REPLACE PARTITION ID '2' FROM \"__temp_target_abcd\"", + "REPLACE PARTITION ID '4' FROM \"__temp_target_abcd\"", + ] + ) + assert sql_calls[3] == 'DROP TABLE IF EXISTS "__temp_target_abcd"' def test_to_time_column():