Skip to content

Commit

Permalink
uses qualified names in queries for test utils
Browse files Browse the repository at this point in the history
  • Loading branch information
rudolfix committed Aug 7, 2023
1 parent 2f9f2ef commit e6b334f
Showing 1 changed file with 13 additions and 9 deletions.
22 changes: 13 additions & 9 deletions tests/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}
Expand All @@ -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}

0 comments on commit e6b334f

Please sign in to comment.