Skip to content

Commit e3f3b7a

Browse files
mday-iomday-io
committed
test: assert ClickHouse overwrite SQL
Signed-off-by: Michael Day <mdaytn@gmail.com>
1 parent 02e6551 commit e3f3b7a

1 file changed

Lines changed: 81 additions & 25 deletions

File tree

‎tests/core/engine_adapter/test_clickhouse.py‎

Lines changed: 81 additions & 25 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,15 +1170,29 @@ def test_insert_overwrite_by_condition_where_partitioned(
11671170
),
11681171
)
11691172

1170-
to_sql_calls(adapter) == [
1173+
sql_calls = to_sql_calls(adapter)
1174+
1175+
assert len(sql_calls) == 7
1176+
assert sql_calls[:5] == [
11711177
'CREATE TABLE "__temp_target_abcd" AS "__temp_existing_records_abcd"',
11721178
"""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")""",
1175-
"""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'""",
1176-
'DROP TABLE IF EXISTS "__temp_target_abcd"',
1179+
'''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"''',
1180+
"""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")""",
1181+
'DROP TABLE IF EXISTS "__temp_affected_partitions_abcd"',
11771182
]
11781183

1184+
alter_prefix = 'ALTER TABLE "__temp_existing_records_abcd" '
1185+
assert sql_calls[5].startswith(alter_prefix)
1186+
assert sorted(sql_calls[5].removeprefix(alter_prefix).split(", ")) == sorted(
1187+
[
1188+
"REPLACE PARTITION ID '1' FROM \"__temp_target_abcd\"",
1189+
"REPLACE PARTITION ID '2' FROM \"__temp_target_abcd\"",
1190+
"REPLACE PARTITION ID '4' FROM \"__temp_target_abcd\"",
1191+
"DROP PARTITION ID '3'",
1192+
]
1193+
)
1194+
assert sql_calls[6] == 'DROP TABLE IF EXISTS "__temp_target_abcd"'
1195+
11791196

11801197
def test_insert_overwrite_by_condition_by_key(
11811198
make_mocked_engine_adapter: t.Callable, mocker: MockerFixture, make_temp_table_name: t.Callable
@@ -1218,15 +1235,15 @@ def test_insert_overwrite_by_condition_by_key(
12181235
dynamic_key_unique=False,
12191236
)
12201237

1221-
to_sql_calls(adapter) == [
1238+
assert to_sql_calls(adapter) == [
12221239
'CREATE TABLE "__temp_target_abcd" AS "__temp_existing_records_abcd"',
12231240
'INSERT INTO "__temp_target_abcd" ("id", "ds") SELECT "id", "ds" FROM (SELECT DISTINCT ON ("id") * FROM "__temp_new_records_abcd") AS "_subquery"',
12241241
'INSERT INTO "__temp_target_abcd" SELECT "id", "ds" FROM "__temp_existing_records_abcd" WHERE NOT ("id" IN (SELECT "id" FROM "__temp_target_abcd"))',
12251242
'EXCHANGE TABLES "__temp_existing_records_abcd" AND "__temp_target_abcd"',
12261243
'DROP TABLE IF EXISTS "__temp_target_abcd"',
12271244
'CREATE TABLE "__temp_target_abcd" AS "__temp_existing_records_abcd"',
12281245
'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"))',
1246+
'INSERT INTO "__temp_target_abcd" SELECT "id", "ds" FROM "__temp_existing_records_abcd" WHERE NOT ("id" IN (SELECT DISTINCT "id" FROM "__temp_target_abcd"))',
12301247
'EXCHANGE TABLES "__temp_existing_records_abcd" AND "__temp_target_abcd"',
12311248
'DROP TABLE IF EXISTS "__temp_target_abcd"',
12321249
]
@@ -1239,17 +1256,20 @@ def test_insert_overwrite_by_condition_by_key_partitioned(
12391256

12401257
temp_table_mock = mocker.patch("sqlmesh.core.engine_adapter.EngineAdapter._get_temp_table")
12411258
table_name = "target"
1242-
temp_table_mock.return_value = make_temp_table_name(table_name, "abcd")
1259+
temp_table_mock.side_effect = lambda table, *args, **kwargs: make_temp_table_name(
1260+
"affected_partitions" if exp.to_table(table).name == "_affected_partitions" else table_name,
1261+
"abcd",
1262+
)
12431263

12441264
fetchone_mock = mocker.patch("sqlmesh.core.engine_adapter.ClickhouseEngineAdapter.fetchone")
12451265
fetchone_mock.side_effect = ["dateTrunc('WEEK', ds)", "dateTrunc('WEEK', ds)"]
12461266

12471267
fetchall_mock = mocker.patch("sqlmesh.core.engine_adapter.ClickhouseEngineAdapter.fetchall")
12481268
fetchall_mock.side_effect = [
12491269
[("1",), ("2",), ("3",), ("4",)],
1250-
["1", "2", "4"],
1270+
[("1",), ("2",), ("4",)],
12511271
[("1",), ("2",), ("3",), ("4",)],
1252-
["1", "2", "4"],
1272+
[("1",), ("2",), ("4",)],
12531273
]
12541274

12551275
insert_table_name = make_temp_table_name("new_records", "abcd")
@@ -1281,20 +1301,45 @@ def test_insert_overwrite_by_condition_by_key_partitioned(
12811301
dynamic_key_unique=False,
12821302
)
12831303

1284-
to_sql_calls(adapter) == [
1304+
sql_calls = to_sql_calls(adapter)
1305+
1306+
assert len(sql_calls) == 14
1307+
assert sql_calls[:4] == [
12851308
'CREATE TABLE "__temp_target_abcd" AS "__temp_existing_records_abcd"',
12861309
'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'""",
1290-
'DROP TABLE IF EXISTS "__temp_target_abcd"',
1310+
'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"',
1311+
'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")',
1312+
]
1313+
1314+
assert sql_calls[4] == 'DROP TABLE IF EXISTS "__temp_affected_partitions_abcd"'
1315+
alter_prefix = 'ALTER TABLE "__temp_existing_records_abcd" '
1316+
assert sql_calls[5].startswith(alter_prefix)
1317+
assert sorted(sql_calls[5].removeprefix(alter_prefix).split(", ")) == sorted(
1318+
[
1319+
"REPLACE PARTITION ID '1' FROM \"__temp_target_abcd\"",
1320+
"REPLACE PARTITION ID '2' FROM \"__temp_target_abcd\"",
1321+
"REPLACE PARTITION ID '4' FROM \"__temp_target_abcd\"",
1322+
"DROP PARTITION ID '3'",
1323+
]
1324+
)
1325+
assert sql_calls[6] == 'DROP TABLE IF EXISTS "__temp_target_abcd"'
1326+
assert sql_calls[7:11] == [
12911327
'CREATE TABLE "__temp_target_abcd" AS "__temp_existing_records_abcd"',
12921328
'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'""",
1296-
'DROP TABLE IF EXISTS "__temp_target_abcd"',
1329+
'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"',
1330+
'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")',
12971331
]
1332+
assert sql_calls[11] == 'DROP TABLE IF EXISTS "__temp_affected_partitions_abcd"'
1333+
assert sql_calls[12].startswith(alter_prefix)
1334+
assert sorted(sql_calls[12].removeprefix(alter_prefix).split(", ")) == sorted(
1335+
[
1336+
"REPLACE PARTITION ID '1' FROM \"__temp_target_abcd\"",
1337+
"REPLACE PARTITION ID '2' FROM \"__temp_target_abcd\"",
1338+
"REPLACE PARTITION ID '4' FROM \"__temp_target_abcd\"",
1339+
"DROP PARTITION ID '3'",
1340+
]
1341+
)
1342+
assert sql_calls[13] == 'DROP TABLE IF EXISTS "__temp_target_abcd"'
12981343

12991344

13001345
def test_insert_overwrite_by_condition_inc_by_partition(
@@ -1330,12 +1375,23 @@ def test_insert_overwrite_by_condition_inc_by_partition(
13301375
keep_existing_partition_rows=False,
13311376
)
13321377

1333-
to_sql_calls(adapter) == [
1378+
sql_calls = to_sql_calls(adapter)
1379+
1380+
assert len(sql_calls) == 4
1381+
assert sql_calls[:2] == [
13341382
'CREATE TABLE "__temp_target_abcd" AS "__temp_existing_records_abcd"',
13351383
'INSERT INTO "__temp_target_abcd" ("id", "ds") SELECT "id", "ds" FROM (SELECT * FROM "__temp_new_records_abcd") AS "_subquery"',
1336-
"""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\"""",
1337-
'DROP TABLE IF EXISTS "__temp_target_abcd"',
13381384
]
1385+
alter_prefix = 'ALTER TABLE "__temp_existing_records_abcd" '
1386+
assert sql_calls[2].startswith(alter_prefix)
1387+
assert sorted(sql_calls[2].removeprefix(alter_prefix).split(", ")) == sorted(
1388+
[
1389+
"REPLACE PARTITION ID '1' FROM \"__temp_target_abcd\"",
1390+
"REPLACE PARTITION ID '2' FROM \"__temp_target_abcd\"",
1391+
"REPLACE PARTITION ID '4' FROM \"__temp_target_abcd\"",
1392+
]
1393+
)
1394+
assert sql_calls[3] == 'DROP TABLE IF EXISTS "__temp_target_abcd"'
13391395

13401396

13411397
def test_to_time_column():

0 commit comments

Comments
 (0)