use metadata.reflect() for DeferredReflection

zzzeek · zzzeek · commit 6cc063f4a217 · 2022-06-21T14:45:25.000-04:00
as metadata.reflect() now has a signficant performance gain over indivdual Table reflection, rework DeferredReflection to use this Fixes: sqlalchemy#8155 Change-Id: I98d55fae83f0cf06a8ca8b89112c85655424d73a
diff --git a/lib/sqlalchemy/ext/declarative/extensions.py b/lib/sqlalchemy/ext/declarative/extensions.py
@@ -10,10 +10,10 @@
 """Public API functions and helpers for declarative."""
 from __future__ import annotations
 
+import collections
 from typing import Callable
 from typing import TYPE_CHECKING
 
-from ... import inspection
 from ...orm import exc as orm_exc
 from ...orm import relationships
 from ...orm.base import _mapper_or_none
@@ -24,7 +24,6 @@
 from ...util import OrderedDict
 
 if TYPE_CHECKING:
-    from ...engine.reflection import Inspector
     from ...sql.schema import MetaData
 
 
@@ -385,12 +384,38 @@ def prepare(cls, engine):
 
         to_map = _DeferredMapperConfig.classes_for_base(cls)
 
-        with inspection.inspect(engine)._inspection_context() as insp:
+        metadata_to_table = collections.defaultdict(set)
+
+        # first collect the primary __table__ for each class into a
+        # collection of metadata/schemaname -> table names
+        for thingy in to_map:
+
+            if thingy.local_table is not None:
+                metadata_to_table[
+                    (thingy.local_table.metadata, thingy.local_table.schema)
+                ].add(thingy.local_table.name)
+
+        # then reflect all those tables into their metadatas
+        with engine.connect() as conn:
+            for (metadata, schema), table_names in metadata_to_table.items():
+                metadata.reflect(
+                    conn,
+                    only=table_names,
+                    schema=schema,
+                    extend_existing=True,
+                    autoload_replace=False,
+                )
+
+            metadata_to_table.clear()
+
+            # .map() each class, then go through relationships and look
+            # for secondary
             for thingy in to_map:
-                cls._sa_decl_prepare(thingy.local_table, insp)
                 thingy.map()
+
                 mapper = thingy.cls.__mapper__
                 metadata = mapper.class_.metadata
+
                 for rel in mapper._props.values():
 
                     if (
@@ -401,41 +426,50 @@ def prepare(cls, engine):
                         secondary_arg = rel._init_args.secondary
 
                         if isinstance(secondary_arg.argument, Table):
-                            cls._reflect_table(secondary_arg.argument, insp)
+                            secondary_table = secondary_arg.argument
+                            metadata_to_table[
+                                (
+                                    secondary_table.metadata,
+                                    secondary_table.schema,
+                                )
+                            ].add(secondary_table.name)
                         elif isinstance(secondary_arg.argument, str):
-
                             _, resolve_arg = _resolver(rel.parent.class_, rel)
 
                             resolver = resolve_arg(
                                 secondary_arg.argument, True
                             )
+                            metadata_to_table[
+                                (metadata, thingy.local_table.schema)
+                            ].add(secondary_arg.argument)
+
                             resolver._resolvers += (
-                                cls._sa_deferred_table_resolver(
-                                    insp, metadata
-                                ),
+                                cls._sa_deferred_table_resolver(metadata),
                             )
 
                             secondary_arg.argument = resolver()
 
+            for (metadata, schema), table_names in metadata_to_table.items():
+                metadata.reflect(
+                    conn,
+                    only=table_names,
+                    schema=schema,
+                    extend_existing=True,
+                    autoload_replace=False,
+                )
+
     @classmethod
     def _sa_deferred_table_resolver(
-        cls, inspector: Inspector, metadata: MetaData
+        cls, metadata: MetaData
     ) -> Callable[[str], Table]:
         def _resolve(key: str) -> Table:
-            t1 = Table(key, metadata)
-            cls._reflect_table(t1, inspector)
-            return t1
+            # reflection has already occurred so this Table would have
+            # its contents already
+            return Table(key, metadata)
 
         return _resolve
 
-    @classmethod
-    def _sa_decl_prepare(cls, local_table, inspector):
-        # autoload Table, which is already
-        # present in the metadata.  This
-        # will fill in db-loaded columns
-        # into the existing Table object.
-        if local_table is not None:
-            cls._reflect_table(local_table, inspector)
+    _sa_decl_prepare = True
 
     @classmethod
     def _sa_raise_deferred_config(cls):
@@ -446,14 +480,3 @@ def _sa_raise_deferred_config(cls):
             "method is called on the class hierarchy."
             % orm_exc._safe_cls_name(cls),
         )
-
-    @classmethod
-    def _reflect_table(cls, table, inspector):
-        Table(
-            table.name,
-            table.metadata,
-            extend_existing=True,
-            autoload_replace=False,
-            autoload_with=inspector,
-            schema=table.schema,
-        )
