From e6b334f09512691f01af89c64d120d7da2adb953 Mon Sep 17 00:00:00 2001 From: Marcin Rudolf Date: Mon, 7 Aug 2023 10:02:44 +0200 Subject: [PATCH] uses qualified names in queries for test utils --- tests/utils.py | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/tests/utils.py b/tests/utils.py index 92f180a9a..77dec67dc 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -196,10 +196,13 @@ def assert_load_info(info: LoadInfo, expected_load_packages: int = 1) -> None: def load_table_counts(p: dlt.Pipeline, *table_names: str) -> DictStrAny: """Returns row counts for `table_names` as dict""" - query = "\nUNION ALL\n".join( - [f"SELECT '{name}' as name, COUNT(1) as c FROM {name}" for name in table_names] - ) with p.sql_client() as c: + query = "\nUNION ALL\n".join( + [ + f"SELECT '{name}' as name, COUNT(1) as c FROM {c.make_qualified_table_name(name)}" + for name in table_names + ] + ) with c.execute_query(query) as cur: rows = list(cur.fetchall()) return {r[0]: r[1] for r in rows} @@ -209,13 +212,14 @@ def load_table_distinct_counts( p: dlt.Pipeline, distinct_column: str, *table_names: str ) -> DictStrAny: """Returns counts of distinct values for column `distinct_column` for `table_names` as dict""" - query = "\nUNION ALL\n".join( - [ - f"SELECT '{name}' as name, COUNT(DISTINCT {distinct_column}) as c FROM {name}" - for name in table_names - ] - ) with p.sql_client() as c: + query = "\nUNION ALL\n".join( + [ + f"SELECT '{name}' as name, COUNT(DISTINCT {distinct_column}) as c FROM {c.make_qualified_table_name(name)}" + for name in table_names + ] + ) + with c.execute_query(query) as cur: rows = list(cur.fetchall()) return {r[0]: r[1] for r in rows}