Tests: Refactor test cases to use generic db helpers

Author: amotl

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()
 

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

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")
 
 

tests/util/db.py

@@ -0,0 +1,32 @@
+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."""
+
+    # CrateDB needs a relaxed SQLAlchemy dialect for querying.
+    # It will not support advanced features of CrateDB,
+    # but that's okay in this case.
+    if uri.startswith("cratedb://"):
+        uri = uri.replace("cratedb://", "postgresql+psycopg_relaxed://")
+
+    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)