new pre-caching architecture for autogenerate

Autogenerate reflection sweeps now use the "bulk" inspector methods
introduced in SQLAlchemy 2.0, which for selected dialects including
PostgreSQL and Oracle use batched queries to reflect whole collections of
tables using O(1) queries rather than O(N).

This is the original proposed version that uses the Inspector
entirely with its public API, with the exception of reflect_table()
which makes a _ReflectionInfo on a per-table basis.  Other than
that, no private API assumptions are made.

If SQLAlchemy needed to add new fields to _ReflectionInfo, it just
needs to make sure they have default functions (which all the current fields
should have anyway, since there is even a ReflectionDefaults
constant that already provides these!)

This version is the one that does not imply any particular
changes in SQLAlchemy and does not have any sqla_compat logic,
so that we may have alembic using the new performance enhancements
allowing for SQLAlchemy to potentially improve its API for a later
release.

Other than that, typing of reflection functions is improved.

Fixes: #1771
Change-Id: I7b9a75fa81cefc377fdb1a22fc1cfc3da1765769

Author: zzzeek

alembic/autogenerate/compare/constraints.py

@@ -5,6 +5,7 @@
 import logging
 from typing import Any
 from typing import cast
+from typing import Collection
 from typing import Dict
 from typing import Mapping
 from typing import Optional
@@ -28,6 +29,9 @@
 from ...util import sqla_compat
 
 if TYPE_CHECKING:
+    from sqlalchemy.engine.interfaces import ReflectedForeignKeyConstraint
+    from sqlalchemy.engine.interfaces import ReflectedIndex
+    from sqlalchemy.engine.interfaces import ReflectedUniqueConstraint
     from sqlalchemy.sql.elements import quoted_name
     from sqlalchemy.sql.elements import TextClause
     from sqlalchemy.sql.schema import Column
@@ -40,7 +44,6 @@
     from ...operations.ops import ModifyTableOps
     from ...runtime.plugins import Plugin
 
-
 _C = TypeVar("_C", bound=Union[UniqueConstraint, ForeignKeyConstraint, Index])
 
 
@@ -72,18 +75,24 @@ def _compare_indexes_and_uniques(
         metadata_unique_constraints = set()
         metadata_indexes = set()
 
-    conn_uniques = conn_indexes = frozenset()  # type:ignore[var-annotated]
+    conn_uniques: Collection[UniqueConstraint] = frozenset()
+    conn_indexes: Collection[Index] = frozenset()
 
     supports_unique_constraints = False
 
     unique_constraints_duplicate_unique_indexes = False
 
     if conn_table is not None:
+        conn_uniques_reflected: Collection[ReflectedUniqueConstraint] = (
+            frozenset()
+        )
+        conn_indexes_reflected: Collection[ReflectedIndex] = frozenset()
+
         # 1b. ... and from connection, if the table exists
         try:
-            conn_uniques = _InspectorConv(inspector).get_unique_constraints(
-                tname, schema=schema
-            )
+            conn_uniques_reflected = _InspectorConv(
+                inspector
+            ).get_unique_constraints(tname, schema=schema)
 
             supports_unique_constraints = True
         except NotImplementedError:
@@ -94,28 +103,28 @@ def _compare_indexes_and_uniques(
             # not being present
             pass
         else:
-            conn_uniques = [  # type:ignore[assignment]
+            conn_uniques_reflected = [
                 uq
-                for uq in conn_uniques
+                for uq in conn_uniques_reflected
                 if autogen_context.run_name_filters(
                     uq["name"],
                     "unique_constraint",
                     {"table_name": tname, "schema_name": schema},
                 )
             ]
-            for uq in conn_uniques:
+            for uq in conn_uniques_reflected:
                 if uq.get("duplicates_index"):
                     unique_constraints_duplicate_unique_indexes = True
         try:
-            conn_indexes = _InspectorConv(inspector).get_indexes(
+            conn_indexes_reflected = _InspectorConv(inspector).get_indexes(
                 tname, schema=schema
             )
         except NotImplementedError:
             pass
         else:
-            conn_indexes = [  # type:ignore[assignment]
+            conn_indexes_reflected = [
                 ix
-                for ix in conn_indexes
+                for ix in conn_indexes_reflected
                 if autogen_context.run_name_filters(
                     ix["name"],
                     "index",
@@ -127,17 +136,18 @@ def _compare_indexes_and_uniques(
         # into schema objects
         if is_drop_table:
             # for DROP TABLE uniques are inline, don't need them
-            conn_uniques = set()  # type:ignore[assignment]
+            conn_uniques = set()
         else:
-            conn_uniques = {  # type:ignore[assignment]
+            conn_uniques = {
                 _make_unique_constraint(impl, uq_def, conn_table)
-                for uq_def in conn_uniques
+                for uq_def in conn_uniques_reflected
             }
 
-        conn_indexes = {  # type:ignore[assignment]
+        conn_indexes = {
             index
             for index in (
-                _make_index(impl, ix, conn_table) for ix in conn_indexes
+                _make_index(impl, ix, conn_table)
+                for ix in conn_indexes_reflected
             )
             if index is not None
         }
@@ -507,7 +517,7 @@ def _correct_for_uq_duplicates_uix(
 
 
 def _make_index(
-    impl: DefaultImpl, params: Dict[str, Any], conn_table: Table
+    impl: DefaultImpl, params: ReflectedIndex, conn_table: Table
 ) -> Optional[Index]:
     exprs: list[Union[Column[Any], TextClause]] = []
     sorting = params.get("column_sorting")
@@ -539,7 +549,7 @@ def _make_index(
 
 
 def _make_unique_constraint(
-    impl: DefaultImpl, params: Dict[str, Any], conn_table: Table
+    impl: DefaultImpl, params: ReflectedUniqueConstraint, conn_table: Table
 ) -> UniqueConstraint:
     uq = sa_schema.UniqueConstraint(
         *[conn_table.c[cname] for cname in params["column_names"]],
@@ -553,7 +563,7 @@ def _make_unique_constraint(
 
 
 def _make_foreign_key(
-    params: Dict[str, Any], conn_table: Table
+    params: ReflectedForeignKeyConstraint, conn_table: Table
 ) -> ForeignKeyConstraint:
     tname = params["referred_table"]
     if params["referred_schema"]:

alembic/autogenerate/compare/tables.py

@@ -48,19 +48,25 @@ def _autogen_for_tables(
     version_table = autogen_context.migration_context.version_table
 
     for schema_name in schemas:
-        tables = set(inspector.get_table_names(schema=schema_name))
+        tables = available = set(inspector.get_table_names(schema=schema_name))
         if schema_name == version_table_schema:
             tables = tables.difference(
                 [autogen_context.migration_context.version_table]
             )
 
-        conn_table_names.update(
-            (schema_name, tname)
+        tablenames = [
+            tname
             for tname in tables
             if autogen_context.run_name_filters(
                 tname, "table", {"schema_name": schema_name}
             )
-        )
+        ]
+
+        conn_table_names.update((schema_name, tname) for tname in tablenames)
+
+        inspector = autogen_context.inspector
+        insp = _InspectorConv(inspector)
+        insp.pre_cache_tables(schema_name, tablenames, available)
 
     metadata_table_names = OrderedSet(
         [(table.schema, table.name) for table in autogen_context.sorted_tables]
@@ -139,6 +145,9 @@ def _compare_tables(
     removal_metadata = sa_schema.MetaData()
     for s, tname in conn_table_names.difference(metadata_table_names):
         name = sa_schema._get_table_key(tname, s)
+
+        # a name might be present already if a previous reflection pulled
+        # this table in via foreign key constraint
         exists = name in removal_metadata.tables
         t = sa_schema.Table(tname, removal_metadata, schema=s)
 
@@ -152,7 +161,7 @@ def _compare_tables(
                 (inspector),
                 # fmt: on
             )
-            _InspectorConv(inspector).reflect_table(t, include_columns=None)
+            _InspectorConv(inspector).reflect_table(t)
         if autogen_context.run_object_filters(t, tname, "table", True, None):
             modify_table_ops = ops.ModifyTableOps(tname, [], schema=s)
 
@@ -172,6 +181,9 @@ def _compare_tables(
     for s, tname in existing_tables:
         name = sa_schema._get_table_key(tname, s)
         exists = name in existing_metadata.tables
+
+        # a name might be present already if a previous reflection pulled
+        # this table in via foreign key constraint
         t = sa_schema.Table(tname, existing_metadata, schema=s)
         if not exists:
             event.listen(
@@ -182,7 +194,7 @@ def _compare_tables(
                 _compat_autogen_column_reflect(inspector),
                 # fmt: on
             )
-            _InspectorConv(inspector).reflect_table(t, include_columns=None)
+            _InspectorConv(inspector).reflect_table(t)
 
         conn_column_info[(s, tname)] = t
 
@@ -296,6 +308,7 @@ def _compare_columns(
 
 
 def setup(plugin: Plugin) -> None:
+
     plugin.add_autogenerate_comparator(
         _autogen_for_tables,
         "schema",