diff --git a/tests/stream/test_kafka.py b/tests/stream/test_kafka.py index a93e85c1..d43f89e0 100644 --- a/tests/stream/test_kafka.py +++ b/tests/stream/test_kafka.py @@ -2,10 +2,10 @@ from concurrent.futures import ThreadPoolExecutor import pytest -import sqlalchemy from confluent_kafka import Producer from tests.util import invoke_ingest_command +from tests.util.db import get_query_result from tests.warehouse.settings import DESTINATIONS # Marked explicitly (not auto-marked by path) because this module lives outside tests/warehouse. @@ -43,13 +43,10 @@ def run(): assert res.exit_code == 0 def get_output_table(): - dest_engine = sqlalchemy.create_engine(dest_uri) - with dest_engine.connect() as conn: - res = conn.exec_driver_sql( - f"select _kafka__data from {topic}.output order by _kafka_msg_id asc" - ).fetchall() - dest_engine.dispose() - return res + return get_query_result( + dest_uri, + f"select _kafka__data from {topic}.output order by _kafka_msg_id asc", + ) run() @@ -118,17 +115,11 @@ def run(): assert res.exit_code == 0 def get_output_table(): - dest_engine = sqlalchemy.create_engine(dest_uri) - with dest_engine.connect() as conn: - res = ( - conn.exec_driver_sql( # ty: ignore[no-matching-overload, unused-ignore-comment, unused-ignore-comment] - f"SELECT id, temperature, humidity FROM {topic}.output WHERE temperature >= 38.00 ORDER BY id ASC" - ) - .mappings() - .fetchall() - ) - dest_engine.dispose() - return res + return get_query_result( + dest_uri, + f"SELECT id, temperature, humidity FROM {topic}.output WHERE temperature >= 38.00 ORDER BY id ASC", + mappings=True, + ) run() @@ -174,17 +165,11 @@ def run(): assert res.exit_code == 0 def get_output_table(): - dest_engine = sqlalchemy.create_engine(dest_uri) - with dest_engine.connect() as conn: - res = ( - conn.exec_driver_sql( - f'SELECT "partition", "topic", "key", "offset" FROM {topic}.output ORDER BY "partition" ASC, "offset" ASC' - ) - .mappings() - .fetchall() - ) - dest_engine.dispose() - return res + return get_query_result( + dest_uri, + f'SELECT "partition", "topic", "key", "offset" FROM {topic}.output ORDER BY "partition" ASC, "offset" ASC', + mappings=True, + ) run() diff --git a/tests/stream/test_mqbridge_kafka.py b/tests/stream/test_mqbridge_kafka.py index 2e34549a..c54b051f 100644 --- a/tests/stream/test_mqbridge_kafka.py +++ b/tests/stream/test_mqbridge_kafka.py @@ -16,10 +16,10 @@ import duckdb import pytest -import sqlalchemy from confluent_kafka import Producer from tests.util import invoke_ingest_command +from tests.util.db import get_query_result from tests.warehouse.settings import DESTINATIONS # Marked explicitly (not auto-marked by path) because this module lives outside tests/warehouse. @@ -64,13 +64,11 @@ def run(): assert res.exit_code == 0 def rows(): - engine = sqlalchemy.create_engine(dest_uri) - with engine.connect() as conn: - out = conn.exec_driver_sql( - f"select order_id, amount from {topic}.output order by order_id asc" - ).fetchall() - engine.dispose() - return [tuple(r) for r in out] + res = get_query_result( + dest_uri, + f"select order_id, amount from {topic}.output order by order_id asc", + ) + return [tuple(r) for r in res] run() assert rows() == EXPECTED diff --git a/tests/util/common.py b/tests/util/common.py index 4f92bc97..e74068b4 100644 --- a/tests/util/common.py +++ b/tests/util/common.py @@ -14,7 +14,7 @@ def as_datetime(date_str: str) -> date: return datetime.strptime(date_str, "%Y-%m-%d").replace(tzinfo=timezone.utc).date() -def as_datetime2(date_str: str) -> datetime: +def as_datetime_notz(date_str: str) -> datetime: return datetime.strptime(date_str, "%Y-%m-%d") diff --git a/tests/util/db.py b/tests/util/db.py new file mode 100644 index 00000000..52620dcf --- /dev/null +++ b/tests/util/db.py @@ -0,0 +1,26 @@ +from typing import Any, Sequence, Union + +import sqlalchemy as sa + + +def dbquery( + uri: str, query: str, fetch: bool = False, mappings: bool = False +) -> Union[Sequence[Union[sa.Row[Any], sa.RowMapping]], None]: + """Query database using SQLAlchemy and optionally return results.""" + + engine = sa.create_engine(uri, poolclass=sa.NullPool) + response = None + with engine.connect() as conn: + res = conn.exec_driver_sql(query) + if fetch: + if mappings: + response = res.mappings().fetchall() + else: + response = res.fetchall() + engine.dispose() + return response + + +def get_query_result(uri: str, query: str, fetch: bool = True, mappings: bool = False): + """Query database using SQLAlchemy and return results.""" + return dbquery(uri, query, fetch=True, mappings=mappings) diff --git a/tests/warehouse/db/test_arrow.py b/tests/warehouse/db/test_arrow.py index ea1d4d12..776d0615 100644 --- a/tests/warehouse/db/test_arrow.py +++ b/tests/warehouse/db/test_arrow.py @@ -6,13 +6,13 @@ import pandas as pd import pyarrow as pa import pytest -import sqlalchemy from pyarrow import ipc as ipc from tests.util import ( invoke_ingest_command, ) -from tests.util.common import as_datetime, as_datetime2, get_random_string +from tests.util.common import as_datetime, as_datetime_notz, get_random_string +from tests.util.db import get_query_result from tests.warehouse.settings import DESTINATIONS @@ -44,7 +44,7 @@ def run_command( ), ) - assert res.exit_code == 0 + assert res.exit_code == 0, res.stderr return res dest_uri = dest.start() @@ -64,17 +64,15 @@ def run_command( table = pa.Table.from_pandas(df) run_command(table) - dest_engine = sqlalchemy.create_engine(dest_uri) - with dest_engine.begin() as conn: - res = conn.exec_driver_sql(f"select count(*) from {schema}.output").fetchall() - assert res[0][0] == row_count + res = get_query_result(dest_uri, f"select count(*) from {schema}.output") + assert res[0][0] == row_count - res = conn.exec_driver_sql( - f"select date, count(*) from {schema}.output group by 1 order by 1 asc" - ).fetchall() - assert res[0][0] == as_datetime("2024-11-05") - assert res[0][1] == row_count - dest_engine.dispose() + res = get_query_result( + dest_uri, + f"select date, count(*) from {schema}.output group by 1 order by 1 asc", + ) + assert res[0][0] == as_datetime("2024-11-05") + assert res[0][1] == row_count # let's add a new column to the dataframe df["new_col"] = "some value" @@ -82,22 +80,22 @@ def run_command( run_command(table) # there should be no change, just a new column - with dest_engine.begin() as conn: - res = conn.exec_driver_sql(f"select count(*) from {schema}.output").fetchall() - assert res[0][0] == row_count + res = get_query_result(dest_uri, f"select count(*) from {schema}.output") + assert res[0][0] == row_count - res = conn.exec_driver_sql( - f"select date, count(*) from {schema}.output group by 1 order by 1 asc" - ).fetchall() - assert res[0][0] == as_datetime("2024-11-05") - assert res[0][1] == row_count + res = get_query_result( + dest_uri, + f"select date, count(*) from {schema}.output group by 1 order by 1 asc", + ) + assert res[0][0] == as_datetime("2024-11-05") + assert res[0][1] == row_count - res = conn.exec_driver_sql( - f"select new_col, count(*) from {schema}.output group by 1 order by 1 asc" - ).fetchall() - assert res[0][0] == "some value" - assert res[0][1] == row_count - dest_engine.dispose() + res = get_query_result( + dest_uri, + f"select new_col, count(*) from {schema}.output group by 1 order by 1 asc", + ) + assert res[0][0] == "some value" + assert res[0][1] == row_count @pytest.mark.parametrize( @@ -123,11 +121,10 @@ def run_command(df: pd.DataFrame, incremental_key: Optional[str] = None): inc_strategy="delete+insert", ) - assert res.exit_code == 0 + assert res.exit_code == 0, res.stderr return res dest_uri = dest.start() - dest_engine = sqlalchemy.create_engine(dest_uri) # let's start with a basic dataframe row_count = 1000 @@ -143,7 +140,7 @@ def run_command(df: pd.DataFrame, incremental_key: Optional[str] = None): run_command(df, "date") def build_datetime(ds: str): - dt: datetime = as_datetime2(ds) + dt: datetime = as_datetime_notz(ds) if dest_uri.startswith("clickhouse"): dt = dt.replace(tzinfo=timezone.utc) return dt @@ -167,29 +164,28 @@ def compare_dates(actual, expected_str): return actual_date == expected_date # the first load, it should be loaded correctly - with dest_engine.begin() as conn: - res = conn.exec_driver_sql(f"select count(*) from {schema}.output").fetchall() - assert res[0][0] == row_count + res = get_query_result(dest_uri, f"select count(*) from {schema}.output") + assert res[0][0] == row_count - res = conn.exec_driver_sql( - f"select date, count(*) from {schema}.output group by 1 order by 1 asc" - ).fetchall() - assert compare_dates(res[0][0], "2024-11-05") - assert res[0][1] == row_count - dest_engine.dispose() + res = get_query_result( + dest_uri, + f"select date, count(*) from {schema}.output group by 1 order by 1 asc", + ) + assert compare_dates(res[0][0], "2024-11-05") + assert res[0][1] == row_count # run again, it should be deleted and reloaded run_command(df, "date") - with dest_engine.begin() as conn: - res = conn.exec_driver_sql(f"select count(*) from {schema}.output").fetchall() - assert res[0][0] == row_count - res = conn.exec_driver_sql( - f"select date, count(*) from {schema}.output group by 1 order by 1 asc" - ).fetchall() - assert compare_dates(res[0][0], "2024-11-05") - assert res[0][1] == row_count - dest_engine.dispose() + res = get_query_result(dest_uri, f"select count(*) from {schema}.output") + assert res[0][0] == row_count + + res = get_query_result( + dest_uri, + f"select date, count(*) from {schema}.output group by 1 order by 1 asc", + ) + assert compare_dates(res[0][0], "2024-11-05") + assert res[0][1] == row_count # append 1000 new rows with a different date new_rows = pd.DataFrame( @@ -204,18 +200,17 @@ def compare_dates(actual, expected_str): run_command(df, "date") - with dest_engine.begin() as conn: - res = conn.exec_driver_sql(f"select count(*) from {schema}.output").fetchall() - assert res[0][0] == row_count + 1000 + res = get_query_result(dest_uri, f"select count(*) from {schema}.output") + assert res[0][0] == row_count + 1000 - res = conn.exec_driver_sql( - f"select date, count(*) from {schema}.output group by 1 order by 1 asc" - ).fetchall() - assert compare_dates(res[0][0], "2024-11-05") - assert res[0][1] == row_count - assert compare_dates(res[1][0], "2024-11-06") - assert res[1][1] == 1000 - dest_engine.dispose() + res = get_query_result( + dest_uri, + f"select date, count(*) from {schema}.output group by 1 order by 1 asc", + ) + assert compare_dates(res[0][0], "2024-11-05") + assert res[0][1] == row_count + assert compare_dates(res[1][0], "2024-11-06") + assert res[1][1] == 1000 # append 1000 old rows for a previous date, these should not be loaded old_rows = pd.DataFrame( @@ -229,18 +224,18 @@ def compare_dates(actual, expected_str): df = pd.concat([df, old_rows], ignore_index=True) run_command(df, "date") - with dest_engine.begin() as conn: - res = conn.exec_driver_sql(f"select count(*) from {schema}.output").fetchall() - assert res[0][0] == row_count + 1000 - res = conn.exec_driver_sql( - f"select date, count(*) from {schema}.output group by 1 order by 1 asc" - ).fetchall() - assert compare_dates(res[0][0], "2024-11-05") - assert res[0][1] == row_count - assert compare_dates(res[1][0], "2024-11-06") - assert res[1][1] == 1000 - dest_engine.dispose() + res = get_query_result(dest_uri, f"select count(*) from {schema}.output") + assert res[0][0] == row_count + 1000 + + res = get_query_result( + dest_uri, + f"select date, count(*) from {schema}.output group by 1 order by 1 asc", + ) + assert compare_dates(res[0][0], "2024-11-05") + assert res[0][1] == row_count + assert compare_dates(res[1][0], "2024-11-06") + assert res[1][1] == 1000 @pytest.mark.parametrize( @@ -265,11 +260,10 @@ def run_command(df: pd.DataFrame): inc_strategy="merge", primary_key="id", ) - assert res.exit_code == 0 + assert res.exit_code == 0, res.stderr return res dest_uri = dest.start() - dest_engine = sqlalchemy.create_engine(dest_uri) # let's start with a basic dataframe row_count = 1000 @@ -278,29 +272,27 @@ def run_command(df: pd.DataFrame): run_command(df) # the first load, it should be loaded correctly - with dest_engine.begin() as conn: - res = conn.exec_driver_sql(f"select count(*) from {schema}.output").fetchall() - assert res[0][0] == row_count + res = get_query_result(dest_uri, f"select count(*) from {schema}.output") + assert res[0][0] == row_count - res = conn.exec_driver_sql( - f"select value, count(*) from {schema}.output group by 1 order by 1 asc" - ).fetchall() - assert res[0][0] == "a" - assert res[0][1] == row_count - dest_engine.dispose() + res = get_query_result( + dest_uri, + f"select value, count(*) from {schema}.output group by 1 order by 1 asc", + ) + assert res[0][0] == "a" + assert res[0][1] == row_count # run again, no change run_command(df) - with dest_engine.begin() as conn: - res = conn.exec_driver_sql(f"select count(*) from {schema}.output").fetchall() - assert res[0][0] == row_count + res = get_query_result(dest_uri, f"select count(*) from {schema}.output") + assert res[0][0] == row_count - res = conn.exec_driver_sql( - f"select value, count(*) from {schema}.output group by 1 order by 1 asc" - ).fetchall() - assert res[0][0] == "a" - assert res[0][1] == row_count - dest_engine.dispose() + res = get_query_result( + dest_uri, + f"select value, count(*) from {schema}.output group by 1 order by 1 asc", + ) + assert res[0][0] == "a" + assert res[0][1] == row_count # append 1000 new rows with a different value new_rows = pd.DataFrame( @@ -313,20 +305,17 @@ def run_command(df: pd.DataFrame): run_command(df) - with dest_engine.begin() as conn: - res = conn.exec_driver_sql(f"select count(*) from {schema}.output").fetchall() - - assert res[0][0] == row_count + 1000 + res = get_query_result(dest_uri, f"select count(*) from {schema}.output") + assert res[0][0] == row_count + 1000 - res = conn.exec_driver_sql( - f"select value, count(*) from {schema}.output group by 1 order by 1 asc" - ).fetchall() - assert res[0][0] == "a" - assert res[0][1] == row_count - assert res[1][0] == "b" - assert res[1][1] == 1000 - - dest_engine.dispose() + res = get_query_result( + dest_uri, + f"select value, count(*) from {schema}.output group by 1 order by 1 asc", + ) + assert res[0][0] == "a" + assert res[0][1] == row_count + assert res[1][0] == "b" + assert res[1][1] == 1000 # append 1000 old rows for previous ids, they should be merged old_rows = pd.DataFrame( @@ -336,12 +325,12 @@ def run_command(df: pd.DataFrame): } ) run_command(old_rows) - with dest_engine.begin() as conn: - res = conn.exec_driver_sql(f"select count(*) from {schema}.output").fetchall() - assert res[0][0] == row_count + 1000 - res = conn.exec_driver_sql( - f"select value, count(*) from {schema}.output group by 1 order by 1 asc" - ).fetchall() - assert res[0][0] == "a" - assert res[0][1] == row_count + 1000 - dest_engine.dispose() + res = get_query_result(dest_uri, f"select count(*) from {schema}.output") + assert res[0][0] == row_count + 1000 + + res = get_query_result( + dest_uri, + f"select value, count(*) from {schema}.output group by 1 order by 1 asc", + ) + assert res[0][0] == "a" + assert res[0][1] == row_count + 1000 diff --git a/tests/warehouse/db/test_couchbase.py b/tests/warehouse/db/test_couchbase.py index 39aacfad..f6a730dd 100644 --- a/tests/warehouse/db/test_couchbase.py +++ b/tests/warehouse/db/test_couchbase.py @@ -3,11 +3,11 @@ from unittest.mock import MagicMock import pytest -import sqlalchemy from omniload.source.couchbase.adapter import fetch_documents from tests.util import invoke_ingest_command from tests.util.container.impl.couchbase import CouchbaseContainer +from tests.util.db import get_query_result from tests.warehouse.manager import COUCHBASE_IMAGE from tests.warehouse.settings import DESTINATIONS @@ -83,13 +83,9 @@ def test_couchbase_source_local(dest): f"Command failed with exit code {result.exit_code}" ) - dest_engine = sqlalchemy.create_engine(dest_uri) - with dest_engine.connect() as conn: - res = conn.exec_driver_sql( - "select * from raw.test_couchbase_collection order by id" - ).fetchall() - dest_engine.dispose() - + res = get_query_result( + dest_uri, "select * from raw.test_couchbase_collection order by id" + ) assert len(res) == 3, f"Expected 3 documents, got {len(res)}" # Verify documents were ingested correctly @@ -170,14 +166,11 @@ def test_couchbase_capella_source(dest): assert result.exit_code == 0, f"Command failed with: {result.output}" # Verify data was ingested - dest_engine = sqlalchemy.create_engine(dest_uri) - with dest_engine.connect() as conn: - res = conn.exec_driver_sql(f"select * from {dest_table}").fetchall() - assert len(res) > 0, "No data was ingested from Couchbase Capella" - print( - f"Successfully ingested {len(res)} documents from Couchbase Capella (bucket in URI)" - ) - dest_engine.dispose() + res = get_query_result(dest_uri, f"select * from {dest_table}") + assert len(res) > 0, "No data was ingested from Couchbase Capella" + print( + f"Successfully ingested {len(res)} documents from Couchbase Capella (bucket in URI)" + ) finally: dest.stop() @@ -222,14 +215,11 @@ def test_couchbase_capella_source_without_bucket_in_uri(dest): assert result.exit_code == 0, f"Command failed with: {result.output}" # Verify data was ingested - dest_engine = sqlalchemy.create_engine(dest_uri) - with dest_engine.connect() as conn: - res = conn.exec_driver_sql(f"select * from {dest_table}").fetchall() - assert len(res) > 0, "No data was ingested from Couchbase Capella" - print( - f"Successfully ingested {len(res)} documents from Couchbase Capella (bucket in table name)" - ) - dest_engine.dispose() + res = get_query_result(dest_uri, f"select * from {dest_table}") + assert len(res) > 0, "No data was ingested from Couchbase Capella" + print( + f"Successfully ingested {len(res)} documents from Couchbase Capella (bucket in table name)" + ) finally: dest.stop() @@ -279,12 +269,9 @@ def test_couchbase_server_source(dest): assert result.exit_code == 0, f"Command failed with: {result.output}" # Verify data was ingested - dest_engine = sqlalchemy.create_engine(dest_uri) - with dest_engine.connect() as conn: - res = conn.exec_driver_sql(f"select * from {dest_table}").fetchall() - assert len(res) > 0, "No data was ingested from Couchbase Server" - print(f"Successfully ingested {len(res)} documents from Couchbase Server") - dest_engine.dispose() + res = get_query_result(dest_uri, f"select * from {dest_table}") + assert len(res) > 0, "No data was ingested from Couchbase Server" + print(f"Successfully ingested {len(res)} documents from Couchbase Server") finally: dest.stop() diff --git a/tests/warehouse/db/test_dynamodb.py b/tests/warehouse/db/test_dynamodb.py index dca3d86f..6e255e58 100644 --- a/tests/warehouse/db/test_dynamodb.py +++ b/tests/warehouse/db/test_dynamodb.py @@ -9,8 +9,8 @@ from tests.util import invoke_ingest_command from tests.util.common import get_random_string from tests.util.container.impl.floci import FlociContainer +from tests.util.db import get_query_result from tests.warehouse.manager import FLOCI_IMAGE -from tests.warehouse.operations import get_query_result from tests.warehouse.settings import DESTINATIONS diff --git a/tests/warehouse/db/test_general.py b/tests/warehouse/db/test_general.py index 8845fe3a..fccc7394 100644 --- a/tests/warehouse/db/test_general.py +++ b/tests/warehouse/db/test_general.py @@ -8,6 +8,7 @@ from tests.util import invoke_ingest_command from tests.util.common import as_datetime, get_random_string from tests.util.container.impl.duckdb import EphemeralDuckDb +from tests.util.db import get_query_result from tests.warehouse.operations import ( custom_query_tests, db_to_db_append, @@ -149,28 +150,26 @@ def test_db_to_db_exclude_columns(source, dest): sql_exclude_columns="col_to_exclude1,col_to_exclude2", ) - assert result.exit_code == 0 + assert result.exit_code == 0, result.stderr - dest_engine = sqlalchemy.create_engine(dest_uri) - with dest_engine.begin() as dest_conn: - res = dest_conn.exec_driver_sql( - f"select id, val, updated_at from {schema_rand_prefix}.output" - ).fetchall() - - assert len(res) == 2 - assert res[0] == (1, "val1", as_datetime("2022-01-01")) - assert res[1] == (2, "val2", as_datetime("2022-02-01")) + res = get_query_result( + dest_uri, + f"select id, val, updated_at from {schema_rand_prefix}.output order by id asc", + ) + assert len(res) == 2 + assert res[0] == (1, "val1", as_datetime("2022-01-01")) + assert res[1] == (2, "val2", as_datetime("2022-02-01")) - # Verify excluded columns don't exist in destination schema - columns = dest_conn.exec_driver_sql( - f"SELECT column_name FROM information_schema.columns " - f"WHERE table_schema = '{schema_rand_prefix}' AND table_name = 'output' " - f"ORDER BY ordinal_position" - ).fetchall() - assert columns == [("id",), ("val",), ("updated_at",)] + # Verify excluded columns don't exist in destination schema + columns = get_query_result( + dest_uri, + f"SELECT column_name FROM information_schema.columns " + f"WHERE table_schema = '{schema_rand_prefix}' AND table_name = 'output' " + f"ORDER BY ordinal_position", + ) + assert columns == [("id",), ("val",), ("updated_at",)] # Clean up - dest_engine.dispose() source.stop() dest.stop() diff --git a/tests/warehouse/db/test_mongodb.py b/tests/warehouse/db/test_mongodb.py index 9c5e4838..acc7d896 100644 --- a/tests/warehouse/db/test_mongodb.py +++ b/tests/warehouse/db/test_mongodb.py @@ -5,12 +5,12 @@ import pandas as pd import pyarrow as pa import pytest -import sqlalchemy from pyarrow import ipc from testcontainers.mongodb import MongoDbContainer from tests.util import invoke_ingest_command from tests.util.common import get_random_string +from tests.util.db import get_query_result from tests.warehouse.manager import MONGODB_IMAGE from tests.warehouse.settings import DESTINATIONS @@ -370,13 +370,10 @@ def test_mongodb_source(mongodb_function, dest): "raw.test_collection", ) - engine = sqlalchemy.create_engine(dest_uri) - with engine.connect() as conn: - res = conn.exec_driver_sql( - "select id, name, nested_parent__key1, nested_parent__key2, nested_parent__key3, key4, value from raw.test_collection order by id" - ).fetchall() - engine.dispose() - + res = get_query_result( + dest_uri, + "select id, name, nested_parent__key1, nested_parent__key2, nested_parent__key3, key4, value from raw.test_collection order by id", + ) assert len(res) == 5 # convert string to json if needed. this is a particular case for Clickhouse which does not have json types by default. @@ -504,12 +501,10 @@ def simple_filtering_query(mongo, dest_uri: str): assert result.exit_code == 0 - engine = sqlalchemy.create_engine(dest_uri) - with engine.connect() as conn: - res = conn.exec_driver_sql( - f"select event_id, event_type, user_id, value from {schema_rand_prefix}.events_success order by event_id" - ).fetchall() - engine.dispose() + res = get_query_result( + dest_uri, + f"select event_id, event_type, user_id, value from {schema_rand_prefix}.events_success order by event_id", + ) assert len(res) == 4 # Only successful events assert res[0] == (1, "login", "user1", 100) @@ -578,12 +573,10 @@ def aggregation_with_grouping(mongo, dest_uri: str): assert result.exit_code == 0 - engine = sqlalchemy.create_engine(dest_uri) - with engine.connect() as conn: - res = conn.exec_driver_sql( - f"select _id, total_value, event_count from {schema_rand_prefix}.user_stats order by _id" - ).fetchall() - engine.dispose() + res = get_query_result( + dest_uri, + f"select _id, total_value, event_count from {schema_rand_prefix}.user_stats order by _id", + ) assert len(res) == 2 assert res[0] == ( @@ -665,12 +658,10 @@ def incremental_with_interval_placeholders(mongo, dest_uri: str): assert result.exit_code == 0 - engine = sqlalchemy.create_engine(dest_uri) - with engine.connect() as conn: - res = conn.exec_driver_sql( - f"select event_id, event_type, user_id, value from {schema_rand_prefix}.events_incremental order by event_id" - ).fetchall() - engine.dispose() + res = get_query_result( + dest_uri, + f"select event_id, event_type, user_id, value from {schema_rand_prefix}.events_incremental order by event_id", + ) # Should only get events from 2024-01-01 (events 1 and 2) assert len(res) == 2 @@ -759,12 +750,10 @@ def incremental_multiple_days(mongo, dest_uri: str): assert result.exit_code == 0 - engine = sqlalchemy.create_engine(dest_uri) - with engine.connect() as conn: - res = conn.exec_driver_sql( - f"select event_id, event_type, user_id, value from {schema_rand_prefix}.events_multi order by event_id" - ).fetchall() - engine.dispose() + res = get_query_result( + dest_uri, + f"select event_id, event_type, user_id, value from {schema_rand_prefix}.events_multi order by event_id", + ) # Should have events from both days (events 1, 2, and 3) assert len(res) == 3 diff --git a/tests/warehouse/db/test_mysql.py b/tests/warehouse/db/test_mysql.py index 0e33e507..e1093112 100644 --- a/tests/warehouse/db/test_mysql.py +++ b/tests/warehouse/db/test_mysql.py @@ -5,6 +5,7 @@ from tests.util import invoke_ingest_command from tests.util.common import get_random_string +from tests.util.db import get_query_result from tests.warehouse.manager import registry from tests.warehouse.settings import DESTINATIONS @@ -57,12 +58,9 @@ def test_mysql_zero_dates(source, dest): assert result.exit_code == 0 - dest_engine = sqlalchemy.create_engine(dest_uri) - with dest_engine.connect() as dest_conn: - res = dest_conn.exec_driver_sql( - f"select * from {schema_rand_prefix}.output order by name" - ).fetchall() - dest_engine.dispose() + res = get_query_result( + dest_uri, f"select * from {schema_rand_prefix}.output order by name" + ) # assert there are no new rows, since DBs like DuckDB accept NULL and dlt adds a separate string column for the value `0000-00-00 00:00:00` # we want 4 columns: name, created_at, _dlt_load_id, _dlt_id diff --git a/tests/warehouse/filesystem/test_local.py b/tests/warehouse/filesystem/test_local.py index bb6a8767..6e57d355 100644 --- a/tests/warehouse/filesystem/test_local.py +++ b/tests/warehouse/filesystem/test_local.py @@ -4,11 +4,11 @@ import pyarrow.csv import pytest -import sqlalchemy from pyarrow import parquet as pya_parquet from tests.util import invoke_ingest_command from tests.util.common import get_random_string, get_testdata_path +from tests.util.db import get_query_result from tests.warehouse.settings import DESTINATIONS TESTDATA = get_testdata_path() @@ -31,12 +31,7 @@ def people_dir(tmp_path_factory): def _scalar(dest_uri, sql): - engine = sqlalchemy.create_engine(dest_uri) - try: - with engine.connect() as conn: - return conn.exec_driver_sql(sql).fetchone() - finally: - engine.dispose() + return get_query_result(dest_uri, sql)[0] @pytest.mark.parametrize( diff --git a/tests/warehouse/filesystem/test_remote.py b/tests/warehouse/filesystem/test_remote.py index a8252e12..e3cee498 100644 --- a/tests/warehouse/filesystem/test_remote.py +++ b/tests/warehouse/filesystem/test_remote.py @@ -9,7 +9,6 @@ import fsspec import pyarrow.csv import pytest -import sqlalchemy from dlt.common.storages.fsspec_filesystem import glob_files from fsspec.implementations.memory import MemoryFileSystem from fsspec.registry import _registry as _fsspec_registry @@ -19,6 +18,7 @@ from omniload.target.filesystem.api import AzureDestination, S3Destination from tests.util import invoke_ingest_command from tests.util.common import get_random_string, has_exception +from tests.util.db import get_query_result from tests.warehouse.settings import DESTINATIONS # These formats need decoders from the optional `iterable` extra. Probe with find_spec so this @@ -122,10 +122,7 @@ def glob_files_override(fs_client, _, file_glob): return glob_files(fs_client, "memory://", file_glob) def assert_rows(dest_uri, dest_table, n): - engine = sqlalchemy.create_engine(dest_uri) - with engine.connect() as conn: - rows = conn.exec_driver_sql(f"select count(*) from {dest_table}").fetchall() - engine.dispose() + rows = get_query_result(dest_uri, f"select count(*) from {dest_table}") assert len(rows) == 1 assert rows[0] == (n,) diff --git a/tests/warehouse/operations.py b/tests/warehouse/operations.py index 8af9275d..447e9f42 100644 --- a/tests/warehouse/operations.py +++ b/tests/warehouse/operations.py @@ -7,6 +7,7 @@ invoke_ingest_command, ) from tests.util.common import as_datetime, get_random_string +from tests.util.db import get_query_result def db_to_db_create_replace(source_connection_url: str, dest_connection_url: str): @@ -38,14 +39,12 @@ def db_to_db_create_replace(source_connection_url: str, dest_connection_url: str f"{schema_rand_prefix}.output", ) - assert result.exit_code == 0 + assert result.exit_code == 0, result.stderr - dest_engine = sqlalchemy.create_engine(dest_connection_url) - with dest_engine.connect() as dest_conn: - res = dest_conn.exec_driver_sql( - f"select id, val, updated_at from {schema_rand_prefix}.output order by id asc" - ).fetchall() - dest_engine.dispose() + res = get_query_result( + dest_connection_url, + f"select id, val, updated_at from {schema_rand_prefix}.output order by id asc", + ) assert len(res) == 2 assert res[0] == (1, "val1", as_datetime("2022-01-01")) @@ -56,7 +55,6 @@ def db_to_db_append(source_connection_url: str, dest_connection_url: str): schema_rand_prefix = f"testschema_append_{get_random_string(5)}" source_engine = sqlalchemy.create_engine(source_connection_url) - dest_engine = sqlalchemy.create_engine(dest_connection_url) with source_engine.begin() as conn: conn.exec_driver_sql(f"DROP SCHEMA IF EXISTS {schema_rand_prefix}") @@ -83,15 +81,13 @@ def run(): "updated_at", sql_backend="sqlalchemy", ) - assert res.exit_code == 0 + assert res.exit_code == 0, res.stderr def get_output_table(): - with dest_engine.connect() as dest_conn: - results = dest_conn.exec_driver_sql( - f"select id, val, updated_at from {schema_rand_prefix}.output order by id asc" - ).fetchall() - dest_engine.dispose() - return results + return get_query_result( + dest_connection_url, + f"select id, val, updated_at from {schema_rand_prefix}.output order by id asc", + ) run() @@ -100,7 +96,7 @@ def get_output_table(): assert res[0] == (1, "val1", as_datetime("2022-01-01")) assert res[1] == (2, "val2", as_datetime("2022-01-02")) - # # run again, nothing should be inserted into the output table + # run again, nothing should be inserted into the output table run() res = get_output_table() @@ -146,16 +142,14 @@ def run(): "id", sql_backend="sqlalchemy", ) - assert res.exit_code == 0 + assert res.exit_code == 0, res.stderr return res - dest_engine = sqlalchemy.create_engine(dest_connection_url) - def get_output_rows(): - with dest_engine.connect() as dest_conn: - return dest_conn.exec_driver_sql( - f"select id, val, updated_at from {schema_rand_prefix}.output order by id asc" - ).fetchall() + return get_query_result( + dest_connection_url, + f"select id, val, updated_at from {schema_rand_prefix}.output order by id asc", + ) def assert_output_equals(expected): res = get_output_rows() @@ -163,18 +157,15 @@ def assert_output_equals(expected): for i, row in enumerate(expected): assert res[i] == row - dest_engine.dispose() res = run() assert_output_equals( [(1, "val1", as_datetime("2022-01-01")), (2, "val2", as_datetime("2022-02-01"))] ) - with dest_engine.connect() as dest_conn: - first_run_id = dest_conn.exec_driver_sql( - f"select _dlt_load_id from {schema_rand_prefix}.output limit 1" - ).fetchall()[0][0] - - dest_engine.dispose() + first_run_id = get_query_result( + dest_connection_url, + f"select _dlt_load_id from {schema_rand_prefix}.output limit 1", + )[0][0] ############################## # we'll run again, we don't expect any changes since the data hasn't changed @@ -184,14 +175,14 @@ def assert_output_equals(expected): ) # we also ensure that the other rows were not touched - with dest_engine.connect() as dest_conn: - count_by_run_id = dest_conn.exec_driver_sql( - f"select _dlt_load_id, count(*) from {schema_rand_prefix}.output group by 1 order by 2 desc" - ).fetchall() + count_by_run_id = get_query_result( + dest_connection_url, + f"select _dlt_load_id, count(*) from {schema_rand_prefix}.output group by 1 order by 2 desc", + ) + assert len(count_by_run_id) == 1 assert count_by_run_id[0][1] == 2 assert count_by_run_id[0][0] == first_run_id - dest_engine.dispose() ############################## ############################## @@ -208,14 +199,13 @@ def assert_output_equals(expected): ) # we also ensure that the other rows were not touched - with dest_engine.connect() as dest_conn: - count_by_run_id = dest_conn.exec_driver_sql( - f"select _dlt_load_id, count(*) from {schema_rand_prefix}.output group by 1" - ).fetchall() + count_by_run_id = get_query_result( + dest_connection_url, + f"select _dlt_load_id, count(*) from {schema_rand_prefix}.output group by 1", + ) assert len(count_by_run_id) == 1 assert count_by_run_id[0][1] == 2 assert count_by_run_id[0][0] == first_run_id - dest_engine.dispose() ############################## ############################## @@ -232,14 +222,13 @@ def assert_output_equals(expected): ) # we also ensure that the other rows were not touched - with dest_engine.connect() as dest_conn: - count_by_run_id = dest_conn.exec_driver_sql( - f"select _dlt_load_id, count(*) from {schema_rand_prefix}.output group by 1" - ).fetchall() + count_by_run_id = get_query_result( + dest_connection_url, + f"select _dlt_load_id, count(*) from {schema_rand_prefix}.output group by 1", + ) assert len(count_by_run_id) == 1 assert count_by_run_id[0][1] == 2 assert count_by_run_id[0][0] == first_run_id - dest_engine.dispose() ############################## ############################## @@ -260,16 +249,15 @@ def assert_output_equals(expected): ) # we have a new run that inserted rows to this table, so the run count should be 2 - with dest_engine.connect() as dest_conn: - count_by_run_id = dest_conn.exec_driver_sql( - f"select _dlt_load_id, count(*) from {schema_rand_prefix}.output group by 1 order by 2 desc" - ).fetchall() + count_by_run_id = get_query_result( + dest_connection_url, + f"select _dlt_load_id, count(*) from {schema_rand_prefix}.output group by 1 order by 2 desc", + ) assert len(count_by_run_id) == 2 assert count_by_run_id[0][1] == 2 assert count_by_run_id[0][0] == first_run_id # we don't care about the run ID assert count_by_run_id[1][1] == 1 - dest_engine.dispose() ############################## ############################## @@ -290,17 +278,16 @@ def assert_output_equals(expected): ) # we have a new run that inserted rows to this table, so the run count should be 2 - with dest_engine.connect() as dest_conn: - count_by_run_id = dest_conn.exec_driver_sql( - f"select _dlt_load_id, count(*) from {schema_rand_prefix}.output group by 1 order by 2 desc, 1 asc" - ).fetchall() + count_by_run_id = get_query_result( + dest_connection_url, + f"select _dlt_load_id, count(*) from {schema_rand_prefix}.output group by 1 order by 2 desc, 1 asc", + ) assert len(count_by_run_id) == 3 assert count_by_run_id[0][1] == 1 assert count_by_run_id[0][0] == first_run_id # we don't care about the rest of the run IDs assert count_by_run_id[1][1] == 1 assert count_by_run_id[2][1] == 1 - dest_engine.dispose() ############################## @@ -473,7 +460,7 @@ def run(start_date: str, end_date: str): interval_end=end_date, sql_backend="sqlalchemy", ) - assert res.exit_code == 0 + assert res.exit_code == 0, res.stderr return res run("2022-01-01", "2022-01-02") # dlt runs them with the end date exclusive @@ -612,14 +599,6 @@ def assert_output_equals(expected): ############################## -def get_query_result(uri: str, query: str): - engine = sqlalchemy.create_engine(uri, poolclass=NullPool) - with engine.connect() as conn: - res = conn.exec_driver_sql(query).fetchall() - engine.dispose() - return res - - def custom_query_tests(): def replace(source_connection_url, dest_connection_url): schema = f"testschema_cr_cust_{get_random_string(5)}" @@ -662,7 +641,7 @@ def replace(source_connection_url, dest_connection_url): run_in_subprocess=True, ) - assert result.exit_code == 0 + assert result.exit_code == 0, result.stderr res = get_query_result( dest_connection_url, @@ -670,7 +649,9 @@ def replace(source_connection_url, dest_connection_url): ) assert len(res) == 4 - assert res[0] == (1, 1, "Item 1 for First Order", as_datetime("2024-01-01")) + assert res[0] == (1, 1, "Item 1 for First Order", as_datetime("2024-01-01")), ( + res[0] + ) assert res[1] == (2, 1, "Item 2 for First Order", as_datetime("2024-01-01")) assert res[2] == (3, 2, "Item 1 for Second Order", as_datetime("2024-01-01")) assert res[3] == (4, 3, "Item 1 for Third Order", as_datetime("2024-01-01")) @@ -713,7 +694,7 @@ def run(): primary_key="id", run_in_subprocess=True, ) - assert result.exit_code == 0 + assert result.exit_code == 0, result.stderr # Initial run to get all data run() diff --git a/tests/warehouse/saas/test_jira.py b/tests/warehouse/saas/test_jira.py index eb85fd9e..0f986695 100644 --- a/tests/warehouse/saas/test_jira.py +++ b/tests/warehouse/saas/test_jira.py @@ -4,10 +4,10 @@ from urllib.parse import quote_plus import pytest -import sqlalchemy from tests.util import invoke_ingest_command from tests.util.common import get_random_string +from tests.util.db import get_query_result from tests.warehouse.settings import DESTINATIONS @@ -60,11 +60,7 @@ def table_test(dest_uri: str): assert result.exit_code == 0 - engine = sqlalchemy.create_engine(dest_uri) - with engine.connect() as conn: - res = conn.exec_driver_sql( - f"select count(*) from {dest_table}" - ).fetchall() + res = get_query_result(dest_uri, f"select count(*) from {dest_table}") assert len(res) > 0, "No results" # Verify the query returned a result count = res[0][0] print(f"Jira {table_name} count: {count}") @@ -77,7 +73,6 @@ def table_test(dest_uri: str): elif table_name == "statuses": assert count > 0, "Jira should have at least one status" # project_versions and project_components can be empty, so no assertion for them - engine.dispose() # Set function name for pytest identification table_test.__name__ = f"{table_name}_table" diff --git a/tests/warehouse/saas/test_pinterest.py b/tests/warehouse/saas/test_pinterest.py index 579fd248..e4990c8f 100644 --- a/tests/warehouse/saas/test_pinterest.py +++ b/tests/warehouse/saas/test_pinterest.py @@ -1,9 +1,9 @@ from unittest.mock import MagicMock, patch import pytest -import sqlalchemy from tests.util import invoke_ingest_command +from tests.util.db import get_query_result from tests.warehouse.settings import DESTINATIONS @@ -93,11 +93,8 @@ def pinterest_test_case(dest_uri): assert result.exit_code == 0 - engine = sqlalchemy.create_engine(dest_uri) - with engine.connect() as conn: - rows = conn.exec_driver_sql(f"SELECT * FROM {dest_table}").fetchall() - assert len(rows) > 0, "No data ingested into the destination" - engine.dispose() + rows = get_query_result(dest_uri, f"SELECT * FROM {dest_table}") + assert len(rows) > 0, "No data ingested into the destination" @pytest.mark.parametrize( diff --git a/tests/warehouse/saas/test_trustpilot.py b/tests/warehouse/saas/test_trustpilot.py index 59d3b5e1..f869c230 100644 --- a/tests/warehouse/saas/test_trustpilot.py +++ b/tests/warehouse/saas/test_trustpilot.py @@ -1,9 +1,9 @@ from unittest.mock import MagicMock, patch import pytest -import sqlalchemy from tests.util import invoke_ingest_command +from tests.util.db import get_query_result from tests.warehouse.settings import DESTINATIONS @@ -115,11 +115,8 @@ def trustpilot_test_case(dest_uri): assert result.exit_code == 0 - engine = sqlalchemy.create_engine(dest_uri) - with engine.connect() as conn: - rows = conn.exec_driver_sql(f"SELECT * FROM {dest_table}").fetchall() - assert len(rows) > 0, "No data ingested into the destination" - engine.dispose() + rows = get_query_result(dest_uri, f"SELECT * FROM {dest_table}") + assert len(rows) > 0, "No data ingested into the destination" @pytest.mark.parametrize(