Merge pull request #45 from DanCardin/dc/agnostic-has-table

DanCardin · web-flow · commit 1bb8552a208f · 2024-01-30T12:36:22.000-05:00
refactor: Use native sqlalchemy dialect function for detecting table …
diff --git a/docs/source/contributing/dialect-support.md b/docs/source/contributing/dialect-support.md
@@ -36,15 +36,12 @@ Here's an excerpt for example:
 ```python
 from sqlalchemy_declarative_extensions.dialects.mysql.query import (
     check_schema_exists_mysql,
-    check_table_exists_mysql,
 )
 from sqlalchemy_declarative_extensions.dialects.postgresql.query import (
     check_schema_exists_postgresql,
-    check_table_exists_postgresql,
 )
 from sqlalchemy_declarative_extensions.dialects.sqlite.query import (
     check_schema_exists_sqlite,
-    check_table_exists_sqlite,
 )
 from sqlalchemy_declarative_extensions.sqlalchemy import dialect_dispatch
 
@@ -53,12 +50,6 @@ check_schema_exists = dialect_dispatch(
     sqlite=check_schema_exists_sqlite,
     mysql=check_schema_exists_mysql,
 )
-
-check_table_exists = dialect_dispatch(
-    postgresql=check_table_exists_postgresql,
-    sqlite=check_table_exists_sqlite,
-    mysql=check_table_exists_mysql,
-)
 ```
 
 "Feature" implementations should only reference **these** functions, which will
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sqlalchemy-declarative-extensions"
-version = "0.6.8"
+version = "0.6.9"
 authors = ["Dan Cardin <ddcardin@gmail.com>"]
 
 description = "Library to declare additional kinds of objects not natively supported by SQLAlchemy/Alembic."
@@ -48,8 +48,8 @@ psycopg = "*"
 alembic-utils = "0.8.1"
 black = ">=22.3.0"
 coverage = ">=5"
-ruff = ">=0.0.165"
-mypy = ">=0.942"
+ruff = "0.0.259"
+mypy = "1.4.1"
 pytest = ">=7"
 pytest-xdist = "*"
 pytest-alembic = "*"
diff --git a/src/sqlalchemy_declarative_extensions/dialects/__init__.py b/src/sqlalchemy_declarative_extensions/dialects/__init__.py
@@ -2,6 +2,7 @@
 from sqlalchemy_declarative_extensions.dialects.query import (
     check_schema_exists,
     check_table_exists,
+    get_current_schema,
     get_default_grants,
     get_function_cls,
     get_functions,
@@ -18,6 +19,7 @@
 __all__ = [
     "check_schema_exists",
     "check_table_exists",
+    "get_current_schema",
     "get_default_grants",
     "get_function_cls",
     "get_grants",
diff --git a/src/sqlalchemy_declarative_extensions/dialects/mysql/query.py b/src/sqlalchemy_declarative_extensions/dialects/mysql/query.py
@@ -2,7 +2,6 @@
 
 from sqlalchemy_declarative_extensions.dialects.mysql.schema import (
     schema_exists_query,
-    table_exists_query,
     views_query,
 )
 from sqlalchemy_declarative_extensions.view.base import View
@@ -23,10 +22,3 @@ def get_views_mysql(connection: Connection):
 def check_schema_exists_mysql(connection: Connection, name: str) -> bool:
     row = connection.execute(schema_exists_query, {"schema": name}).scalar()
     return not bool(row)
-
-
-def check_table_exists_mysql(connection: Connection, name: str, *, schema: str) -> bool:
-    row = connection.execute(
-        table_exists_query, {"name": name, "schema": schema}
-    ).scalar()
-    return bool(row)
diff --git a/src/sqlalchemy_declarative_extensions/dialects/mysql/schema.py b/src/sqlalchemy_declarative_extensions/dialects/mysql/schema.py
@@ -33,12 +33,6 @@
     .where(views.c.table_schema.notin_(["sys"]))
 )
 
-table_exists_query = (
-    select(tables)
-    .where(tables.c.table_schema == bindparam("schema"))
-    .where(tables.c.table_name == bindparam("name"))
-)
-
 schema_exists_query = select(schemata).where(
     schemata.c.schema_name == bindparam("schema")
 )
diff --git a/src/sqlalchemy_declarative_extensions/dialects/postgresql/query.py b/src/sqlalchemy_declarative_extensions/dialects/postgresql/query.py
@@ -18,7 +18,6 @@
     roles_query,
     schema_exists_query,
     schemas_query,
-    table_exists_query,
     triggers_query,
     view_query,
     views_query,
@@ -46,15 +45,6 @@ def check_schema_exists_postgresql(connection: Connection, name: str) -> bool:
     return not bool(row)
 
 
-def check_table_exists_postgresql(
-    connection: Connection, name: str, *, schema: str
-) -> bool:
-    row = connection.execute(
-        table_exists_query, {"name": name, "schema": schema}
-    ).scalar()
-    return bool(row)
-
-
 def get_objects_postgresql(connection: Connection):
     return sorted(
         [
diff --git a/src/sqlalchemy_declarative_extensions/dialects/postgresql/schema.py b/src/sqlalchemy_declarative_extensions/dialects/postgresql/schema.py
@@ -150,13 +150,6 @@ def _schema_not_pg(column=pg_namespace.c.nspname):
 )
 
 
-table_exists_query = (
-    select(tables)
-    .where(tables.c.table_schema == bindparam("schema"))
-    .where(tables.c.table_name == bindparam("name"))
-)
-
-
 default_acl_query = select(
     pg_roles.c.rolname.label("role_name"),
     pg_namespace.c.nspname.label("schema_name"),
diff --git a/src/sqlalchemy_declarative_extensions/dialects/query.py b/src/sqlalchemy_declarative_extensions/dialects/query.py
@@ -1,12 +1,15 @@
+from __future__ import annotations
+
+from sqlalchemy import func
+from sqlalchemy.engine import Connection
+
 from sqlalchemy_declarative_extensions.dialects import postgresql
 from sqlalchemy_declarative_extensions.dialects.mysql.query import (
     check_schema_exists_mysql,
-    check_table_exists_mysql,
     get_views_mysql,
 )
 from sqlalchemy_declarative_extensions.dialects.postgresql.query import (
     check_schema_exists_postgresql,
-    check_table_exists_postgresql,
     get_default_grants_postgresql,
     get_functions_postgresql,
     get_grants_postgresql,
@@ -19,10 +22,9 @@
 )
 from sqlalchemy_declarative_extensions.dialects.sqlite.query import (
     check_schema_exists_sqlite,
-    check_table_exists_sqlite,
     get_views_sqlite,
 )
-from sqlalchemy_declarative_extensions.sqlalchemy import dialect_dispatch
+from sqlalchemy_declarative_extensions.sqlalchemy import dialect_dispatch, select
 
 get_schemas = dialect_dispatch(
     postgresql=get_schemas_postgresql,
@@ -34,12 +36,6 @@
     mysql=check_schema_exists_mysql,
 )
 
-check_table_exists = dialect_dispatch(
-    postgresql=check_table_exists_postgresql,
-    sqlite=check_table_exists_sqlite,
-    mysql=check_table_exists_mysql,
-)
-
 get_objects = dialect_dispatch(
     postgresql=get_objects_postgresql,
 )
@@ -85,3 +81,22 @@
 get_triggers = dialect_dispatch(
     postgresql=get_triggers_postgresql,
 )
+
+
+def check_table_exists(
+    connection: Connection, tablename: str, schema: str | None = None
+):
+    return connection.dialect.has_table(connection, tablename, schema=schema)
+
+
+def get_current_schema(connection: Connection) -> str | None:
+    if connection.dialect.name == "mysql":
+        return None
+
+    schema = connection.execute(select(func.current_schema())).scalar()
+
+    default_schema = connection.dialect.default_schema_name
+    if schema == default_schema:
+        return None
+
+    return schema.lower()
diff --git a/src/sqlalchemy_declarative_extensions/dialects/sqlite/query.py b/src/sqlalchemy_declarative_extensions/dialects/sqlite/query.py
@@ -1,10 +1,7 @@
 from sqlalchemy import text
 from sqlalchemy.engine import Connection
 
-from sqlalchemy_declarative_extensions.dialects.sqlite.schema import (
-    table_exists_query,
-    views_query,
-)
+from sqlalchemy_declarative_extensions.dialects.sqlite.schema import views_query
 from sqlalchemy_declarative_extensions.view.base import View
 
 
@@ -26,13 +23,3 @@ def get_views_sqlite(connection: Connection):
         View(v.name, v.definition, schema=v.schema)
         for v in connection.execute(views_query()).fetchall()
     ]
-
-
-def check_table_exists_sqlite(
-    connection: Connection, name: str, *, schema: str
-) -> bool:
-    row = connection.execute(
-        table_exists_query(schema),
-        {"name": name, "schema": schema},
-    ).scalar()
-    return bool(row)
diff --git a/src/sqlalchemy_declarative_extensions/dialects/sqlite/schema.py b/src/sqlalchemy_declarative_extensions/dialects/sqlite/schema.py
@@ -1,6 +1,6 @@
 from typing import Optional
 
-from sqlalchemy import bindparam, column, literal, table
+from sqlalchemy import column, literal, table
 
 from sqlalchemy_declarative_extensions.sqlalchemy import select
 
@@ -25,13 +25,3 @@ def views_query(schema: Optional[str] = None):
         sqlite_schema.c.name.label("name"),
         sqlite_schema.c.sql.label("definition"),
     ).where(sqlite_schema.c.type == "view")
-
-
-def table_exists_query(schema: Optional[str] = None):
-    sqlite_schema = make_sqlite_schema(schema)
-    return (
-        select(sqlite_schema)
-        .where(sqlite_schema.c.type == "table")
-        .where(sqlite_schema.c.table_schema == bindparam("schema"))
-        .where(sqlite_schema.c.name == bindparam("name"))
-    )
diff --git a/src/sqlalchemy_declarative_extensions/row/__init__.py b/src/sqlalchemy_declarative_extensions/row/__init__.py
@@ -1,8 +1,9 @@
 from sqlalchemy_declarative_extensions.row import compare
-from sqlalchemy_declarative_extensions.row.base import Row, Rows
+from sqlalchemy_declarative_extensions.row.base import Row, Rows, Table
 
 __all__ = [
     "compare",
     "Rows",
     "Row",
+    "Table",
 ]
diff --git a/src/sqlalchemy_declarative_extensions/row/base.py b/src/sqlalchemy_declarative_extensions/row/base.py
@@ -33,15 +33,15 @@ class Row:
     tablename: str
     column_values: dict[str, Any]
 
-    def __init__(self, tablename, **column_values):
+    def __init__(self, tablename, *, schema: str | None = None, **column_values):
+        schema = schema
         try:
             schema, table = tablename.split(".", 1)
         except ValueError:
-            self.schema = None
-            self.tablename = tablename
-        else:
-            self.schema = schema
-            self.tablename = table
+            table = tablename
+
+        self.schema = schema
+        self.tablename = table
 
         self.column_values = column_values
 
@@ -51,6 +51,19 @@ def qualified_name(self):
             return f"{self.schema}.{self.tablename}"
         return self.tablename
 
+    def qualify(self, schema: str | None) -> Row:
+        """Attach a schema to this row, if it doesn't have one already.
+
+        Examples:
+            >>> row = Row("foo")
+            >>> row.qualify("bar")
+            Row(schema='bar', tablename='foo', column_values={})
+        """
+        if self.schema is not None:
+            return self
+
+        return self.__class__(self.tablename, schema=schema, **self.column_values)
+
 
 @dataclass
 class Table:
@@ -74,6 +87,7 @@ class Table:
         ... ]
         [Row(schema=None, tablename='users', column_values={'active': True, 'id': 1, 'name': 'John'}), Row(schema=None, tablename='users', column_values={'active': True, 'id': 2, 'name': 'Bob'})]
     """
+
     name: str
     column_values: dict[str, Any]
 
@@ -82,5 +96,5 @@ def __init__(self, name, **column_values):
         self.column_values = column_values
 
     def row(self, **column_values) -> Row:
-        final_values= {**self.column_values, **column_values}
+        final_values = {**self.column_values, **column_values}
         return Row(self.name, **final_values)
diff --git a/src/sqlalchemy_declarative_extensions/row/compare.py b/src/sqlalchemy_declarative_extensions/row/compare.py
@@ -6,7 +6,10 @@
 from sqlalchemy import MetaData, Table, tuple_
 from sqlalchemy.engine.base import Connection
 
-from sqlalchemy_declarative_extensions.dialects import check_table_exists
+from sqlalchemy_declarative_extensions.dialects import (
+    check_table_exists,
+    get_current_schema,
+)
 from sqlalchemy_declarative_extensions.row.base import Row, Rows
 from sqlalchemy_declarative_extensions.sqlalchemy import row_to_dict
 
@@ -100,8 +103,12 @@ def get_table(conn: Connection, tablename: str):
 def compare_rows(connection: Connection, metadata: MetaData, rows: Rows) -> list[RowOp]:
     result: list[RowOp] = []
 
+    current_schema = get_current_schema(connection)
+
     rows_by_table: dict[Table, list[Row]] = {}
     for row in rows:
+        row = row.qualify(current_schema)
+
         table = metadata.tables.get(row.qualified_name)
         if table is None:
             raise ValueError(f"Unknown table: {row.qualified_name}")
@@ -124,7 +131,7 @@ def compare_rows(connection: Connection, metadata: MetaData, rows: Rows) -> list
         table_exists = check_table_exists(
             connection,
             table.name,
-            schema=table.schema or connection.dialect.default_schema_name,
+            schema=table.schema or current_schema,
         )
 
         record = None
@@ -160,7 +167,7 @@ def compare_rows(connection: Connection, metadata: MetaData, rows: Rows) -> list
             table_exists = check_table_exists(
                 connection,
                 table.name,
-                schema=table.schema or connection.dialect.default_schema_name,
+                schema=table.schema or current_schema,
             )
             if not table_exists:
                 continue
diff --git a/src/sqlalchemy_declarative_extensions/sqlalchemy.py b/src/sqlalchemy_declarative_extensions/sqlalchemy.py
@@ -39,10 +39,7 @@ def escape_params(query: str) -> str:
 
 
 if version.startswith("1.3"):
-    from sqlalchemy.ext.declarative import (
-        DeclarativeMeta,
-        instrument_declarative,
-    )
+    from sqlalchemy.ext.declarative import DeclarativeMeta, instrument_declarative
 
     def select(*args):
         return sqlalchemy.select(list(args))
diff --git a/src/sqlalchemy_declarative_extensions/view/base.py b/src/sqlalchemy_declarative_extensions/view/base.py
@@ -39,11 +39,11 @@ def view(
     as though it were a table.
 
     Arguments:
-        base: A declarative base object.
+        base: A declarative base object
         materialized: Whether the view should be a materialized view
         register_as_model: Whether the view should be registered as a SQLAlchemy mapped object.
             Note this only works if the view defines mappable models columns (minimally a primary
-            key), like a proper modeled table.
+            key), like a proper modeled table
 
     >>> try:
     ...     from sqlalchemy.orm import declarative_base

Original file line number	Diff line number	Diff line change
`@@ -33,12 +33,6 @@`
`33`	`33`	`.where(views.c.table_schema.notin_(["sys"]))`
`34`	`34`	`)`
`35`	`35`
`36`		`-table_exists_query = (`
`37`		`- select(tables)`
`38`		`- .where(tables.c.table_schema == bindparam("schema"))`
`39`		`- .where(tables.c.table_name == bindparam("name"))`
`40`		`-)`
`41`		`-`
`42`	`36`	`schema_exists_query = select(schemata).where(`
`43`	`37`	`schemata.c.schema_name == bindparam("schema")`
`44`	`38`	`)`