chore: Simplify type annotations for SQLAlchemy v2 (#17)

Author: borchero

sqlcompyre/analysis/dialects/_base.py

@@ -1,11 +1,10 @@
-# Copyright (c) QuantCo 2024-2024
+# Copyright (c) QuantCo 2024-2025
 # SPDX-License-Identifier: BSD-3-Clause
 
 from datetime import datetime
 from typing import Protocol
 
 import sqlalchemy as sa
-from sqlalchemy.engine import Engine
 
 
 class DialectProtocol(Protocol):
@@ -40,7 +39,7 @@ class DialectProtocol(Protocol):
     views_support_notnull_columns: bool
 
     def get_table_creation_timestamps(
-        self, engine: Engine, tables: list[sa.Table]
+        self, engine: sa.Engine, tables: list[sa.Table]
     ) -> list[datetime]:
         """Obtain the creation timestamps from a list of tables.
 

sqlcompyre/analysis/dialects/mssql.py

@@ -1,11 +1,10 @@
-# Copyright (c) QuantCo 2024-2024
+# Copyright (c) QuantCo 2024-2025
 # SPDX-License-Identifier: BSD-3-Clause
 
 from datetime import datetime
 
 import sqlalchemy as sa
 from sqlalchemy.dialects.mssql import dialect as SqlAlchemyMssqlDialect  # noqa: N812
-from sqlalchemy.engine import Engine
 
 from ._base import DialectProtocol
 
@@ -22,7 +21,7 @@ class MssqlDialect(SqlAlchemyMssqlDialect, DialectProtocol):  # type: ignore
     views_support_notnull_columns: bool = True
 
     def get_table_creation_timestamps(
-        self, engine: Engine, tables: list[sa.Table]
+        self, engine: sa.Engine, tables: list[sa.Table]
     ) -> list[datetime]:
         # Potentially, we need to get the database from the tables
         db: str | None = None

sqlcompyre/analysis/query_inspection.py

@@ -1,11 +1,10 @@
-# Copyright (c) QuantCo 2024-2024
+# Copyright (c) QuantCo 2024-2025
 # SPDX-License-Identifier: BSD-3-Clause
 
 
 from functools import cached_property, lru_cache
 
 import sqlalchemy as sa
-from sqlalchemy.engine import Engine
 
 
 class QueryInspection:
@@ -16,7 +15,7 @@ class QueryInspection:
         or :meth:`~sqlcompyre.api.inspect_table` functions instead.
     """
 
-    def __init__(self, engine: Engine, selectable: sa.Select):
+    def __init__(self, engine: sa.Engine, selectable: sa.Select):
         """
         Args:
             engine: The engine to use for connecting to the database.

sqlcompyre/analysis/schema_comparison.py

@@ -1,13 +1,12 @@
-# Copyright (c) QuantCo 2024-2024
+# Copyright (c) QuantCo 2024-2025
 # SPDX-License-Identifier: BSD-3-Clause
 
 import logging
 import re
 from functools import cached_property
 from typing import Literal, cast
 
-from sqlalchemy import schema
-from sqlalchemy.engine import Engine
+import sqlalchemy as sa
 from tqdm.auto import tqdm
 
 from sqlcompyre.report import Report
@@ -27,11 +26,11 @@ class SchemaComparison:
 
     def __init__(
         self,
-        engine: Engine,
+        engine: sa.Engine,
         left_schema: str,
         right_schema: str,
-        left_tables: dict[str, schema.Table],
-        right_tables: dict[str, schema.Table],
+        left_tables: dict[str, sa.Table],
+        right_tables: dict[str, sa.Table],
         float_precision: float,
         collation: str | None,
         ignore_casing: bool,

sqlcompyre/analysis/table_comparison.py

@@ -6,9 +6,6 @@
 from functools import cached_property
 
 import sqlalchemy as sa
-import sqlalchemy.sql.functions as func
-from sqlalchemy.engine import Engine
-from sqlalchemy.sql import elements, expression, false, select, selectable, true
 
 from sqlcompyre.report import Report
 from sqlcompyre.results import ColumnMatches, Counts, Names, RowMatches
@@ -24,7 +21,7 @@ class TableComparison:
 
     def __init__(
         self,
-        engine: Engine,
+        engine: sa.Engine,
         left_table: sa.FromClause,
         right_table: sa.FromClause,
         join_columns: list[str] | None,
@@ -169,14 +166,14 @@ def row_matches(self) -> RowMatches:
             for colname_1, colname_2 in self.column_name_mapping.items()
             if colname_1 not in self.join_columns
         ]
-        inequality_conditions: list[elements.ColumnElement[bool]] = [
+        inequality_conditions: list[sa.ColumnElement[bool]] = [
             sa.not_(c) for c in equality_conditions
         ]
 
         # If there are no conditions, equality is always true, inequality is always false
         if not equality_conditions:
-            equality_conditions = [true()]
-            inequality_conditions = [false()]
+            equality_conditions = [sa.true()]
+            inequality_conditions = [sa.false()]
 
         # -- Create queries
         # Query for rows ONLY in left table
@@ -186,7 +183,7 @@ def row_matches(self) -> RowMatches:
             if c not in self.join_columns
         ]
         unjoined_left = (
-            select(*left_columns)
+            sa.select(*left_columns)
             .select_from(self._outer_join(left=True))
             .where(
                 self.right_table.c[self.column_name_mapping[self.join_columns[0]]].is_(
@@ -204,7 +201,7 @@ def row_matches(self) -> RowMatches:
             if k not in self.join_columns
         ]
         unjoined_right = (
-            select(*right_columns)
+            sa.select(*right_columns)
             .select_from(self._outer_join(left=False))
             .where(self.left_table.c[self.join_columns[0]].is_(None))
         )
@@ -229,7 +226,7 @@ def row_matches(self) -> RowMatches:
         ]
 
         # The remaining queries
-        joined_total = select(*columns).select_from(self._inner_join())
+        joined_total = sa.select(*columns).select_from(self._inner_join())
         joined_unequal = joined_total.where(sa.or_(*inequality_conditions))
         joined_equal = joined_total.where(sa.and_(*equality_conditions))
         joined_row_count = self._count_rows(self._inner_join())
@@ -266,11 +263,11 @@ def column_matches(self) -> ColumnMatches:
         if len(cases) == 0:
             return ColumnMatches(fraction_same={}, mismatch_selects={})
 
-        case_stmt = select(*cases).select_from(inner_join).subquery()
+        case_stmt = sa.select(*cases).select_from(inner_join).subquery()
 
         # Compute fraction of matching values
         cols_to_avg = [col for col in case_stmt.c if f"_{MATCH_SUFFIX}" in col.name]
-        avgs = select(
+        avgs = sa.select(
             *[
                 sa.func.avg(col).label(f"{col.name.replace(f'_{MATCH_SUFFIX}', '')}")
                 for col in cols_to_avg
@@ -284,7 +281,7 @@ def column_matches(self) -> ColumnMatches:
 
         # Find column mismatches
         mismatch_selects = {
-            left_column: select(inner_join).where(
+            left_column: sa.select(inner_join).where(
                 sa.not_(self._is_equal(left_column, right_column))
             )
             for left_column, right_column in self.column_name_mapping.items()
@@ -374,9 +371,7 @@ def _right_table_name(self) -> str:
             return str(self.right_table.element)
         return "<right query>"
 
-    def _is_equal(
-        self, left_column: str, right_column: str
-    ) -> elements.ColumnElement[bool]:
+    def _is_equal(self, left_column: str, right_column: str) -> sa.ColumnElement[bool]:
         """Forms a condition for comparing two columns.
 
         Args:
@@ -403,13 +398,13 @@ def _is_equal(
         # and inverting this is still `unknown`). For more discussion, see
         # https://stackoverflow.com/questions/1075142/how-to-compare-values-which-may-both-be-null-in-t-sql
         # The following is a more robust formulation of `A = B OR (A IS NULL AND B IS NULL)`.
-        return func.coalesce(
+        return sa.func.coalesce(
             sa.case((condition, None), else_=lhs),
             sa.case((condition, None), else_=rhs),
         ).is_(None)
 
     @cached_property
-    def _join_conditions(self) -> list[elements.ColumnElement[bool]]:
+    def _join_conditions(self) -> list[sa.ColumnElement[bool]]:
         """Forms a list of join conditions."""
         return [
             (
@@ -419,11 +414,11 @@ def _join_conditions(self) -> list[elements.ColumnElement[bool]]:
             for join_col in self.join_columns
         ]
 
-    def _inner_join(self) -> expression.Join:
+    def _inner_join(self) -> sa.Join:
         """Specifies an inner join on the left and right tables."""
         return self.left_table.join(self.right_table, sa.and_(*self._join_conditions))
 
-    def _outer_join(self, left: bool) -> expression.Join:
+    def _outer_join(self, left: bool) -> sa.Join:
         """Specifies an outer join between the two tables.
 
         Args:
@@ -438,7 +433,7 @@ def _outer_join(self, left: bool) -> expression.Join:
             return left_table.outerjoin(right_table, sa.and_(*self._join_conditions))
         return right_table.outerjoin(left_table, sa.and_(*self._join_conditions))
 
-    def _get_aggregate_changes(self, left_col_name: str) -> selectable.Select:
+    def _get_aggregate_changes(self, left_col_name: str) -> sa.Select:
         """Counts the number of different ways each column changes from one table to
         another.
 
@@ -462,14 +457,14 @@ def _get_aggregate_changes(self, left_col_name: str) -> selectable.Select:
         )
 
         return (
-            select(change, sa.func.count())
+            sa.select(change, sa.func.count())
             .select_from(self._inner_join())
             .where(sa.not_(self._is_equal(left_col_name, right_col_name)))
             .group_by(left_col, right_col)
             .order_by(sa.func.count().desc())
         )
 
-    def _count_rows(self, table: expression.FromClause) -> int:
+    def _count_rows(self, table: sa.FromClause) -> int:
         """Counts the number of rows in a table-like object.
 
         Args:
@@ -479,7 +474,9 @@ def _count_rows(self, table: expression.FromClause) -> int:
             The number of rows.
         """
         with self.engine.connect() as conn:
-            return conn.execute(select(sa.func.count()).select_from(table)).scalar_one()
+            return conn.execute(
+                sa.select(sa.func.count()).select_from(table)
+            ).scalar_one()
 
     # ---------------------------------------------------------------------------------------------
     # STRING REPRESENTATION
@@ -502,7 +499,7 @@ def __str__(self):
 
 
 def _join_columns_from_pk_if_needed(
-    engine: Engine,
+    engine: sa.Engine,
     left: sa.FromClause,
     right: sa.FromClause,
     join_columns: list[str],
@@ -564,7 +561,7 @@ def _join_columns_from_pk_if_needed(
 
 
 def _is_valid_primary_key_column(
-    engine: Engine,
+    engine: sa.Engine,
     left_table: sa.FromClause,
     right_table: sa.FromClause,
     left_column: str,
@@ -593,7 +590,7 @@ def _is_valid_primary_key_column(
 
 
 def _is_valid_primary_key(
-    engine: Engine, table: sa.FromClause, columns: list[str]
+    engine: sa.Engine, table: sa.FromClause, columns: list[str]
 ) -> bool:
     with engine.connect() as conn:
         result = conn.execute(

sqlcompyre/api.py

@@ -4,7 +4,6 @@
 import sys
 
 import sqlalchemy as sa
-from sqlalchemy import schema
 
 from .analysis import QueryInspection, SchemaComparison, TableComparison
 
@@ -248,7 +247,7 @@ def _get_tables_from_schema(
     schema: str,
     is_database: bool,
     include_views: bool,
-) -> list[schema.Table]:
+) -> list[sa.Table]:
     if is_database:
         engine = sa.create_engine(engine.url.set(database=schema))
         schemas = sa.inspect(engine).get_schema_names()

sqlcompyre/results/column_matches.py

@@ -1,9 +1,9 @@
-# Copyright (c) QuantCo 2024-2024
+# Copyright (c) QuantCo 2024-2025
 # SPDX-License-Identifier: BSD-3-Clause
 
 from dataclasses import dataclass
 
-from sqlalchemy.sql import selectable
+import sqlalchemy as sa
 
 
 @dataclass
@@ -15,4 +15,4 @@ class ColumnMatches:
     fraction_same: dict[str, float]
     #: Dictionary mapping the name of the left-table column to a query of all joined rows for
     #: which the column does not have the same value in both tables.
-    mismatch_selects: dict[str, selectable.Select]
+    mismatch_selects: dict[str, sa.Select]

sqlcompyre/results/row_matches.py

@@ -1,9 +1,9 @@
-# Copyright (c) QuantCo 2024-2024
+# Copyright (c) QuantCo 2024-2025
 # SPDX-License-Identifier: BSD-3-Clause
 
 from dataclasses import dataclass
 
-from sqlalchemy.sql import selectable
+import sqlalchemy as sa
 
 
 @dataclass
@@ -21,12 +21,12 @@ class RowMatches:
     #: Number of rows that could be joined
     n_joined_total: int
     #: Query for obtaining all rows from the left table that could not be joined.
-    unjoined_left: selectable.Select
+    unjoined_left: sa.Select
     #: Query for obtaining all rows from the right table that could not be joined.
-    unjoined_right: selectable.Select
+    unjoined_right: sa.Select
     #: Query for obtaining all rows that could be joined and were identical across the two tables.
-    joined_equal: selectable.Select
+    joined_equal: sa.Select
     #: Query for obtaining all rows that could be joined but were not identical.
-    joined_unequal: selectable.Select
+    joined_unequal: sa.Select
     #: Query for obtaining all rows that were joined, regardless of equality.
-    joined_total: selectable.Select
+    joined_total: sa.Select

tests/_shared/dialects.py

@@ -1,11 +1,11 @@
-# Copyright (c) QuantCo 2024-2024
+# Copyright (c) QuantCo 2024-2025
 # SPDX-License-Identifier: BSD-3-Clause
 
 from __future__ import annotations
 
 import os
 
-from sqlalchemy.engine import url
+import sqlalchemy as sa
 
 from sqlcompyre.analysis.dialects import DialectProtocol
 
@@ -17,15 +17,13 @@ def dialect_from_env() -> DialectProtocol:
     connection string. **Use this function with care!** Typically, it should only be
     used for ``pytest.mark.skipif`` statements.
     """
-    conn = url.make_url(os.environ["DB_CONNECTION_STRING"])
+    conn = sa.make_url(os.environ["DB_CONNECTION_STRING"])
     return dialect_from_connection_url(conn)
 
 
-def dialect_from_connection_url(conn_url: url.URL) -> DialectProtocol:
+def dialect_from_connection_url(conn_url: sa.URL) -> DialectProtocol:
     """Get dialect metadata from the connection URL."""
-    from sqlalchemy import create_engine
-
-    dialect = create_engine(conn_url).dialect
+    dialect = sa.create_engine(conn_url).dialect
     match dialect.name:
         case "mssql":
             from sqlcompyre.analysis.dialects import MssqlDialect