Skip to content

Commit 3c00cb4

Browse files
kuephiphilipp-kuebler
kuephi
and
philipp-kuebler
authored
Use a make_insert_query function in test_sql.py (#3983)
* added make_insert_query function * used make_insert_query function in first occurance of load_data_query * added types and docstring * refactored all remanining places to use the new make_insert_query function * edited docstring to please linter --------- Co-authored-by: kuephi <philipp.kuebler@aoe.com>
1 parent abfefc0 commit 3c00cb4

File tree

2 files changed

+40
-67
lines changed

2 files changed

+40
-67
lines changed

‎catalog/tests/dags/common/loader/test_sql.py

Lines changed: 34 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -447,9 +447,7 @@ def test_upsert_records_inserts_one_record_to_empty_image_table(
447447
col.HEIGHT.db_name: HEIGHT,
448448
}
449449
)
450-
load_data_query = f"""INSERT INTO {load_table} VALUES(
451-
{query_values}
452-
);"""
450+
load_data_query = utils.make_insert_query(load_table, query_values)
453451

454452
_set_up_std_popularity_func(
455453
postgres_with_load_and_image_table,
@@ -686,9 +684,8 @@ def test_upsert_records_replaces_data(
686684
col.HEIGHT.db_name: HEIGHT_A,
687685
}
688686
)
689-
load_data_query_a = f"""INSERT INTO {load_table} VALUES(
690-
{query_values}
691-
);"""
687+
load_data_query_a = utils.make_insert_query(load_table, query_values)
688+
692689
_set_up_std_popularity_func(
693690
postgres_with_load_and_image_table,
694691
load_data_query_a,
@@ -728,9 +725,8 @@ def test_upsert_records_replaces_data(
728725
col.HEIGHT.db_name: HEIGHT_B,
729726
}
730727
)
731-
load_data_query_b = f"""INSERT INTO {load_table} VALUES(
732-
{query_values}
733-
);"""
728+
load_data_query_b = utils.make_insert_query(load_table, query_values)
729+
734730
postgres_with_load_and_image_table.cursor.execute(f"DELETE FROM {load_table};")
735731
postgres_with_load_and_image_table.connection.commit()
736732
postgres_with_load_and_image_table.cursor.execute(load_data_query_b)
@@ -817,9 +813,8 @@ def test_upsert_records_does_not_replace_with_nulls(
817813
col.HEIGHT.db_name: HEIGHT_A,
818814
}
819815
)
820-
load_data_query_a = f"""INSERT INTO {load_table} VALUES(
821-
{query_values_a}
822-
);"""
816+
load_data_query_a = utils.make_insert_query(load_table, query_values_a)
817+
823818
_set_up_std_popularity_func(
824819
postgres_with_load_and_image_table,
825820
load_data_query_a,
@@ -849,9 +844,8 @@ def test_upsert_records_does_not_replace_with_nulls(
849844
col.SOURCE.db_name: SOURCE,
850845
}
851846
)
852-
load_data_query_b = f"""INSERT INTO {load_table} VALUES(
853-
{query_values_b}
854-
);"""
847+
load_data_query_b = utils.make_insert_query(load_table, query_values_b)
848+
855849
postgres_with_load_and_image_table.cursor.execute(f"DELETE FROM {load_table};")
856850
postgres_with_load_and_image_table.connection.commit()
857851
postgres_with_load_and_image_table.cursor.execute(load_data_query_b)
@@ -910,9 +904,7 @@ def test_upsert_records_merges_meta_data(
910904
col.PROVIDER.db_name: PROVIDER,
911905
}
912906
)
913-
load_data_query_a = f"""INSERT INTO {load_table} VALUES(
914-
{query_values_a}
915-
);"""
907+
load_data_query_a = utils.make_insert_query(load_table, query_values_a)
916908

917909
query_values_b = utils.create_query_values(
918910
{
@@ -923,9 +915,8 @@ def test_upsert_records_merges_meta_data(
923915
col.PROVIDER.db_name: PROVIDER,
924916
}
925917
)
926-
load_data_query_b = f"""INSERT INTO {load_table} VALUES(
927-
{query_values_b}
928-
);"""
918+
load_data_query_b = utils.make_insert_query(load_table, query_values_b)
919+
929920
_set_up_std_popularity_func(
930921
postgres_with_load_and_image_table,
931922
load_data_query_a,
@@ -991,9 +982,7 @@ def test_upsert_records_does_not_replace_with_null_values_in_meta_data(
991982
col.PROVIDER.db_name: PROVIDER,
992983
}
993984
)
994-
load_data_query_a = f"""INSERT INTO {load_table} VALUES(
995-
{query_values_a}
996-
);"""
985+
load_data_query_a = utils.make_insert_query(load_table, query_values_a)
997986

998987
query_values_b = utils.create_query_values(
999988
{
@@ -1004,9 +993,8 @@ def test_upsert_records_does_not_replace_with_null_values_in_meta_data(
1004993
col.PROVIDER.db_name: PROVIDER,
1005994
}
1006995
)
1007-
load_data_query_b = f"""INSERT INTO {load_table} VALUES(
1008-
{query_values_b}
1009-
);"""
996+
load_data_query_b = utils.make_insert_query(load_table, query_values_b)
997+
1010998
_set_up_std_popularity_func(
1011999
postgres_with_load_and_image_table,
10121000
load_data_query_a,
@@ -1081,9 +1069,7 @@ def test_upsert_records_merges_tags(
10811069
col.PROVIDER.db_name: PROVIDER,
10821070
}
10831071
)
1084-
load_data_query_a = f"""INSERT INTO {load_table} VALUES(
1085-
{query_values_a}
1086-
);"""
1072+
load_data_query_a = utils.make_insert_query(load_table, query_values_a)
10871073

10881074
query_values_b = utils.create_query_values(
10891075
{
@@ -1094,9 +1080,7 @@ def test_upsert_records_merges_tags(
10941080
col.PROVIDER.db_name: PROVIDER,
10951081
}
10961082
)
1097-
load_data_query_b = f"""INSERT INTO {load_table} VALUES(
1098-
{query_values_b}
1099-
);"""
1083+
load_data_query_b = utils.make_insert_query(load_table, query_values_b)
11001084
_set_up_std_popularity_func(
11011085
postgres_with_load_and_image_table,
11021086
load_data_query_a,
@@ -1169,9 +1153,7 @@ def test_upsert_records_does_not_replace_tags_with_null(
11691153
col.PROVIDER.db_name: PROVIDER,
11701154
}
11711155
)
1172-
load_data_query_a = f"""INSERT INTO {load_table} VALUES(
1173-
{query_values_a}
1174-
);"""
1156+
load_data_query_a = utils.make_insert_query(load_table, query_values_a)
11751157

11761158
query_values_b = utils.create_query_values(
11771159
{
@@ -1181,9 +1163,8 @@ def test_upsert_records_does_not_replace_tags_with_null(
11811163
col.PROVIDER.db_name: PROVIDER,
11821164
}
11831165
)
1184-
load_data_query_b = f"""INSERT INTO {load_table} VALUES(
1185-
{query_values_b}
1186-
);"""
1166+
load_data_query_b = utils.make_insert_query(load_table, query_values_b)
1167+
11871168
_set_up_std_popularity_func(
11881169
postgres_with_load_and_image_table,
11891170
load_data_query_a,
@@ -1253,9 +1234,8 @@ def test_upsert_records_replaces_null_tags(
12531234
}
12541235
)
12551236
logging.info(f"Query values a: {query_values_a}")
1256-
load_data_query_a = f"""INSERT INTO {load_table} VALUES(
1257-
{query_values_a}
1258-
);"""
1237+
load_data_query_a = utils.make_insert_query(load_table, query_values_a)
1238+
12591239
query_values_b = utils.create_query_values(
12601240
{
12611241
col.FOREIGN_ID.db_name: FID,
@@ -1265,9 +1245,7 @@ def test_upsert_records_replaces_null_tags(
12651245
col.PROVIDER.db_name: PROVIDER,
12661246
}
12671247
)
1268-
load_data_query_b = f"""INSERT INTO {load_table} VALUES(
1269-
{query_values_b}
1270-
);"""
1248+
load_data_query_b = utils.make_insert_query(load_table, query_values_b)
12711249

12721250
_set_up_std_popularity_func(
12731251
postgres_with_load_and_image_table,
@@ -1342,9 +1320,7 @@ def test_upsert_records_handles_duplicate_url_and_does_not_merge(
13421320
col.PROVIDER.db_name: PROVIDER,
13431321
}
13441322
)
1345-
load_data_query_a = f"""INSERT INTO {load_table} VALUES(
1346-
{query_values_a}
1347-
);"""
1323+
load_data_query_a = utils.make_insert_query(load_table, query_values_a)
13481324

13491325
query_values_b = utils.create_query_values(
13501326
{
@@ -1355,9 +1331,7 @@ def test_upsert_records_handles_duplicate_url_and_does_not_merge(
13551331
col.PROVIDER.db_name: PROVIDER,
13561332
}
13571333
)
1358-
load_data_query_b = f"""INSERT INTO {load_table} VALUES(
1359-
{query_values_b}
1360-
);"""
1334+
load_data_query_b = utils.make_insert_query(load_table, query_values_b)
13611335

13621336
# Simulate a DAG run where A is ingested into the loading table, upserted into
13631337
# the image table, and finally the loading table is cleared for the next DAG run.
@@ -1438,9 +1412,7 @@ def test_upsert_records_handles_duplicate_urls_in_a_single_batch_and_does_not_me
14381412
col.PROVIDER.db_name: PROVIDER,
14391413
}
14401414
)
1441-
load_data_query_a = f"""INSERT INTO {load_table} VALUES(
1442-
{query_values_a}
1443-
);"""
1415+
load_data_query_a = utils.make_insert_query(load_table, query_values_a)
14441416

14451417
query_values_b = utils.create_query_values(
14461418
{
@@ -1451,9 +1423,7 @@ def test_upsert_records_handles_duplicate_urls_in_a_single_batch_and_does_not_me
14511423
col.PROVIDER.db_name: PROVIDER,
14521424
}
14531425
)
1454-
load_data_query_b = f"""INSERT INTO {load_table} VALUES(
1455-
{query_values_b}
1456-
);"""
1426+
load_data_query_b = utils.make_insert_query(load_table, query_values_b)
14571427

14581428
# C is not a duplicate of anything, just a normal image
14591429
query_values_c = utils.create_query_values(
@@ -1465,9 +1435,7 @@ def test_upsert_records_handles_duplicate_urls_in_a_single_batch_and_does_not_me
14651435
col.PROVIDER.db_name: PROVIDER,
14661436
}
14671437
)
1468-
load_data_query_c = f"""INSERT INTO {load_table} VALUES(
1469-
{query_values_c}
1470-
);"""
1438+
load_data_query_c = utils.make_insert_query(load_table, query_values_c)
14711439

14721440
# Simulate a DAG run where duplicates (A and B) and a non-duplicate (C) are all
14731441
# ingested in a single batch from the provider script, and we attempt to upsert
@@ -1595,13 +1563,13 @@ def test_upsert_records_calculates_standardized_popularity(
15951563
)
15961564

15971565
# Queries to insert the two records into the load table.
1598-
load_data_query_old_record = f"""INSERT INTO {load_table} VALUES(
1599-
{query_values_update_old_record}
1600-
);"""
1566+
load_data_query_old_record = utils.make_insert_query(
1567+
load_table, query_values_update_old_record
1568+
)
16011569

1602-
load_data_query_new_record = f"""INSERT INTO {load_table} VALUES(
1603-
{query_values_update_new_record}
1604-
);"""
1570+
load_data_query_new_record = utils.make_insert_query(
1571+
load_table, query_values_update_new_record
1572+
)
16051573

16061574
# Now actually insert the records into the load table. This simulates a DagRun which ingests both
16071575
# records.

‎catalog/tests/test_utils/sql.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -109,13 +109,18 @@ def create_query_values(
109109
return ",".join(result)
110110

111111

112+
def make_insert_query(table: str, values: str) -> str:
113+
"""Return an SQL insert statement for the given table with the given values"""
114+
return f"INSERT INTO {table} VALUES({values});"
115+
116+
112117
def _get_insert_query(image_table, values: dict):
113118
# Append the required identifier
114119
values[col.IDENTIFIER.db_name] = uuid.uuid4()
115120

116121
query_values = create_query_values(values, columns=IMAGE_TABLE_COLUMNS)
117122

118-
return f"INSERT INTO {image_table} VALUES({query_values});"
123+
return make_insert_query(image_table, query_values)
119124

120125

121126
def load_sample_data_into_image_table(image_table, postgres, records):

0 commit comments

Comments
 (0)