Skip to content

Commit 966e34d

Browse files
mday-iomday-io
authored
test: assert ClickHouse overwrite SQL (#6149)
Restore effective ClickHouse overwrite SQL assertions and sort partition actions for deterministic SQL. Signed-off-by: Michael Day <mdaytn@gmail.com>
1 parent a8e3405 commit 966e34d

2 files changed

Lines changed: 30 additions & 21 deletions

File tree

‎sqlmesh/core/engine_adapter/clickhouse.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -446,7 +446,7 @@ def _build_alter_partition_exp(
446446
) -> exp.Alter:
447447
alter_expr = exp.Alter(this=target_table, kind="TABLE")
448448

449-
for partition in partitions_to_replace:
449+
for partition in sorted(partitions_to_replace):
450450
alter_expr.append(
451451
"actions",
452452
exp.ReplacePartition(
@@ -457,7 +457,7 @@ def _build_alter_partition_exp(
457457
),
458458
)
459459

460-
for partition in partitions_to_drop:
460+
for partition in sorted(partitions_to_drop):
461461
alter_expr.append(
462462
"actions",
463463
exp.DropPartition(

‎tests/core/engine_adapter/test_clickhouse.py‎

Lines changed: 28 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1118,7 +1118,7 @@ def test_insert_overwrite_by_condition_replace(
11181118
columns_to_types,
11191119
)
11201120

1121-
to_sql_calls(adapter) == [
1121+
assert to_sql_calls(adapter) == [
11221122
'CREATE TABLE "__temp_target_abcd" AS "__temp_existing_records_abcd"',
11231123
'INSERT INTO "__temp_target_abcd" ("id", "ds") SELECT "id", "ds" FROM (SELECT * FROM "__temp_new_records_abcd") AS "_subquery"',
11241124
'EXCHANGE TABLES "__temp_existing_records_abcd" AND "__temp_target_abcd"',
@@ -1133,15 +1133,18 @@ def test_insert_overwrite_by_condition_where_partitioned(
11331133

11341134
temp_table_mock = mocker.patch("sqlmesh.core.engine_adapter.EngineAdapter._get_temp_table")
11351135
table_name = "target"
1136-
temp_table_mock.return_value = make_temp_table_name(table_name, "abcd")
1136+
temp_table_mock.side_effect = lambda table, *args, **kwargs: make_temp_table_name(
1137+
"affected_partitions" if exp.to_table(table).name == "_affected_partitions" else table_name,
1138+
"abcd",
1139+
)
11371140

11381141
fetchone_mock = mocker.patch("sqlmesh.core.engine_adapter.ClickhouseEngineAdapter.fetchone")
11391142
fetchone_mock.return_value = "dateTrunc('WEEK', ds)"
11401143

11411144
fetchall_mock = mocker.patch("sqlmesh.core.engine_adapter.ClickhouseEngineAdapter.fetchall")
11421145
fetchall_mock.side_effect = [
11431146
[("1",), ("2",), ("3",), ("4",)],
1144-
["1", "2", "4"],
1147+
[("1",), ("2",), ("4",)],
11451148
]
11461149

11471150
insert_table_name = make_temp_table_name("new_records", "abcd")
@@ -1167,11 +1170,12 @@ def test_insert_overwrite_by_condition_where_partitioned(
11671170
),
11681171
)
11691172

1170-
to_sql_calls(adapter) == [
1173+
assert to_sql_calls(adapter) == [
11711174
'CREATE TABLE "__temp_target_abcd" AS "__temp_existing_records_abcd"',
11721175
"""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'""",
1173-
"""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\"""",
1174-
"""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")""",
1176+
'''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"''',
1177+
"""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")""",
1178+
'DROP TABLE IF EXISTS "__temp_affected_partitions_abcd"',
11751179
"""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'""",
11761180
'DROP TABLE IF EXISTS "__temp_target_abcd"',
11771181
]
@@ -1218,15 +1222,15 @@ def test_insert_overwrite_by_condition_by_key(
12181222
dynamic_key_unique=False,
12191223
)
12201224

1221-
to_sql_calls(adapter) == [
1225+
assert to_sql_calls(adapter) == [
12221226
'CREATE TABLE "__temp_target_abcd" AS "__temp_existing_records_abcd"',
12231227
'INSERT INTO "__temp_target_abcd" ("id", "ds") SELECT "id", "ds" FROM (SELECT DISTINCT ON ("id") * FROM "__temp_new_records_abcd") AS "_subquery"',
12241228
'INSERT INTO "__temp_target_abcd" SELECT "id", "ds" FROM "__temp_existing_records_abcd" WHERE NOT ("id" IN (SELECT "id" FROM "__temp_target_abcd"))',
12251229
'EXCHANGE TABLES "__temp_existing_records_abcd" AND "__temp_target_abcd"',
12261230
'DROP TABLE IF EXISTS "__temp_target_abcd"',
12271231
'CREATE TABLE "__temp_target_abcd" AS "__temp_existing_records_abcd"',
12281232
'INSERT INTO "__temp_target_abcd" ("id", "ds") SELECT "id", "ds" FROM (SELECT * FROM "__temp_new_records_abcd") AS "_subquery"',
1229-
'INSERT INTO "__temp_target_abcd" SELECT "id", "ds" FROM "__temp_existing_records_abcd" WHERE NOT ("id" IN (SELECT "id" FROM "__temp_target_abcd"))',
1233+
'INSERT INTO "__temp_target_abcd" SELECT "id", "ds" FROM "__temp_existing_records_abcd" WHERE NOT ("id" IN (SELECT DISTINCT "id" FROM "__temp_target_abcd"))',
12301234
'EXCHANGE TABLES "__temp_existing_records_abcd" AND "__temp_target_abcd"',
12311235
'DROP TABLE IF EXISTS "__temp_target_abcd"',
12321236
]
@@ -1239,17 +1243,20 @@ def test_insert_overwrite_by_condition_by_key_partitioned(
12391243

12401244
temp_table_mock = mocker.patch("sqlmesh.core.engine_adapter.EngineAdapter._get_temp_table")
12411245
table_name = "target"
1242-
temp_table_mock.return_value = make_temp_table_name(table_name, "abcd")
1246+
temp_table_mock.side_effect = lambda table, *args, **kwargs: make_temp_table_name(
1247+
"affected_partitions" if exp.to_table(table).name == "_affected_partitions" else table_name,
1248+
"abcd",
1249+
)
12431250

12441251
fetchone_mock = mocker.patch("sqlmesh.core.engine_adapter.ClickhouseEngineAdapter.fetchone")
12451252
fetchone_mock.side_effect = ["dateTrunc('WEEK', ds)", "dateTrunc('WEEK', ds)"]
12461253

12471254
fetchall_mock = mocker.patch("sqlmesh.core.engine_adapter.ClickhouseEngineAdapter.fetchall")
12481255
fetchall_mock.side_effect = [
12491256
[("1",), ("2",), ("3",), ("4",)],
1250-
["1", "2", "4"],
1257+
[("1",), ("2",), ("4",)],
12511258
[("1",), ("2",), ("3",), ("4",)],
1252-
["1", "2", "4"],
1259+
[("1",), ("2",), ("4",)],
12531260
]
12541261

12551262
insert_table_name = make_temp_table_name("new_records", "abcd")
@@ -1281,18 +1288,20 @@ def test_insert_overwrite_by_condition_by_key_partitioned(
12811288
dynamic_key_unique=False,
12821289
)
12831290

1284-
to_sql_calls(adapter) == [
1291+
assert to_sql_calls(adapter) == [
12851292
'CREATE TABLE "__temp_target_abcd" AS "__temp_existing_records_abcd"',
12861293
'INSERT INTO "__temp_target_abcd" ("id", "ds") SELECT "id", "ds" FROM (SELECT DISTINCT ON ("id") * FROM "__temp_new_records_abcd") AS "_subquery"',
1287-
'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"',
1288-
'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")',
1289-
"""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'""",
1294+
'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"',
1295+
'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")',
1296+
'DROP TABLE IF EXISTS "__temp_affected_partitions_abcd"',
1297+
"""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'""",
12901298
'DROP TABLE IF EXISTS "__temp_target_abcd"',
12911299
'CREATE TABLE "__temp_target_abcd" AS "__temp_existing_records_abcd"',
12921300
'INSERT INTO "__temp_target_abcd" ("id", "ds") SELECT "id", "ds" FROM (SELECT * FROM "__temp_new_records_abcd") AS "_subquery"',
1293-
'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"',
1294-
'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")',
1295-
"""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'""",
1301+
'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"',
1302+
'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")',
1303+
'DROP TABLE IF EXISTS "__temp_affected_partitions_abcd"',
1304+
"""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'""",
12961305
'DROP TABLE IF EXISTS "__temp_target_abcd"',
12971306
]
12981307

@@ -1330,7 +1339,7 @@ def test_insert_overwrite_by_condition_inc_by_partition(
13301339
keep_existing_partition_rows=False,
13311340
)
13321341

1333-
to_sql_calls(adapter) == [
1342+
assert to_sql_calls(adapter) == [
13341343
'CREATE TABLE "__temp_target_abcd" AS "__temp_existing_records_abcd"',
13351344
'INSERT INTO "__temp_target_abcd" ("id", "ds") SELECT "id", "ds" FROM (SELECT * FROM "__temp_new_records_abcd") AS "_subquery"',
13361345
"""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\"""",

0 commit comments

Comments
 (0)