apply criteria options from top-level core-only statement

zzzeek · zzzeek · commit a05ae2c7ce0c · 2023-04-17T15:32:00.000-04:00
Made an improvement to the :func:`_orm.with_loader_criteria` loader option to allow it to be indicated in the :meth:`.Executable.options` method of a top-level statement that is not itself an ORM statement. Examples include :func:`_sql.select` that's embedded in compound statements such as :func:`_sql.union`, within an :meth:`_dml.Insert.from_select` construct, as well as within CTE expressions that are not ORM related at the top level. Improved propagation of :func:`_orm.with_loader_criteria` within ORM enabled UPDATE and DELETE statements as well. Fixes: sqlalchemy#9635 Change-Id: I088ad91929dc797c06f292f5dc547d48ffb30430
diff --git a/doc/build/changelog/unreleased_20/9635.rst b/doc/build/changelog/unreleased_20/9635.rst
@@ -0,0 +1,10 @@
+.. change::
+    :tags: bug, orm
+    :tickets: 9635
+
+    Made an improvement to the :func:`_orm.with_loader_criteria` loader option
+    to allow it to be indicated in the :meth:`.Executable.options` method of a
+    top-level statement that is not itself an ORM statement. Examples include
+    :func:`_sql.select` that's embedded in compound statements such as
+    :func:`_sql.union`, within an :meth:`_dml.Insert.from_select` construct, as
+    well as within CTE expressions that are not ORM related at the top level.
diff --git a/lib/sqlalchemy/orm/bulk_persistence.py b/lib/sqlalchemy/orm/bulk_persistence.py
@@ -1346,15 +1346,14 @@ def _setup_for_orm_update(self, statement, compiler, **kw):
 
         self.mapper = mapper = ext_info.mapper
 
-        self.extra_criteria_entities = {}
-
         self._resolved_values = self._get_resolved_values(mapper, statement)
 
-        extra_criteria_attributes = {}
-
-        for opt in statement._with_options:
-            if opt._is_criteria_option:
-                opt.get_global_criteria(extra_criteria_attributes)
+        self._init_global_attributes(
+            statement,
+            compiler,
+            toplevel=True,
+            process_criteria_for_toplevel=True,
+        )
 
         if statement._values:
             self._resolved_values = dict(self._resolved_values)
@@ -1372,7 +1371,7 @@ def _setup_for_orm_update(self, statement, compiler, **kw):
             new_stmt._values = self._resolved_values
 
         new_crit = self._adjust_for_extra_criteria(
-            extra_criteria_attributes, mapper
+            self.global_attributes, mapper
         )
         if new_crit:
             new_stmt = new_stmt.where(*new_crit)
@@ -1741,19 +1740,18 @@ def create_for_statement(cls, statement, compiler, **kw):
         ext_info = statement.table._annotations["parententity"]
         self.mapper = mapper = ext_info.mapper
 
-        self.extra_criteria_entities = {}
-
-        extra_criteria_attributes = {}
-
-        for opt in statement._with_options:
-            if opt._is_criteria_option:
-                opt.get_global_criteria(extra_criteria_attributes)
+        self._init_global_attributes(
+            statement,
+            compiler,
+            toplevel=True,
+            process_criteria_for_toplevel=True,
+        )
 
         new_stmt = statement._clone()
         new_stmt.table = mapper.local_table
 
         new_crit = cls._adjust_for_extra_criteria(
-            extra_criteria_attributes, mapper
+            self.global_attributes, mapper
         )
         if new_crit:
             new_stmt = new_stmt.where(*new_crit)
diff --git a/lib/sqlalchemy/orm/context.py b/lib/sqlalchemy/orm/context.py
@@ -209,6 +209,45 @@ def _get_top_level_context(self) -> QueryContext:
 class AbstractORMCompileState(CompileState):
     is_dml_returning = False
 
+    def _init_global_attributes(
+        self, statement, compiler, *, toplevel, process_criteria_for_toplevel
+    ):
+        self.attributes = {}
+
+        if compiler is None:
+            # this is the legacy / testing only ORM _compile_state() use case.
+            # there is no need to apply criteria options for this.
+            self.global_attributes = ga = {}
+            assert toplevel
+            return
+        else:
+            self.global_attributes = ga = compiler._global_attributes
+
+        if toplevel:
+            ga["toplevel_orm"] = True
+
+            if process_criteria_for_toplevel:
+                for opt in statement._with_options:
+                    if opt._is_criteria_option:
+                        opt.process_compile_state(self)
+
+            return
+        elif ga.get("toplevel_orm", False):
+            return
+
+        stack_0 = compiler.stack[0]
+
+        try:
+            toplevel_stmt = stack_0["selectable"]
+        except KeyError:
+            pass
+        else:
+            for opt in toplevel_stmt._with_options:
+                if opt._is_compile_state and opt._is_criteria_option:
+                    opt.process_compile_state(self)
+
+        ga["toplevel_orm"] = True
+
     @classmethod
     def create_for_statement(
         cls,
@@ -622,17 +661,13 @@ def create_for_statement(
 
         assert isinstance(statement_container, FromStatement)
 
-        if compiler is not None:
-            toplevel = not compiler.stack
-        else:
-            toplevel = True
-
-        if not toplevel:
+        if compiler is not None and compiler.stack:
             raise sa_exc.CompileError(
                 "The ORM FromStatement construct only supports being "
                 "invoked as the topmost statement, as it is only intended to "
                 "define how result rows should be returned."
             )
+
         self = cls.__new__(cls)
         self._primary_entity = None
 
@@ -680,18 +715,18 @@ def create_for_statement(
 
         self.current_path = statement_container._compile_options._current_path
 
-        if toplevel and statement_container._with_options:
-            self.attributes = {}
-            self.global_attributes = compiler._global_attributes
+        self._init_global_attributes(
+            statement_container,
+            compiler,
+            process_criteria_for_toplevel=False,
+            toplevel=True,
+        )
 
+        if statement_container._with_options:
             for opt in statement_container._with_options:
                 if opt._is_compile_state:
                     opt.process_compile_state(self)
 
-        else:
-            self.attributes = {}
-            self.global_attributes = compiler._global_attributes
-
         if statement_container._with_context_options:
             for fn, key in statement_container._with_context_options:
                 fn(self)
@@ -911,10 +946,8 @@ def create_for_statement(
 
         if compiler is not None:
             toplevel = not compiler.stack
-            self.global_attributes = compiler._global_attributes
         else:
             toplevel = True
-            self.global_attributes = {}
 
         select_statement = statement
 
@@ -1002,11 +1035,17 @@ def create_for_statement(
 
         self.eager_order_by = ()
 
+        self._init_global_attributes(
+            select_statement,
+            compiler,
+            toplevel=toplevel,
+            process_criteria_for_toplevel=False,
+        )
+
         if toplevel and (
             select_statement._with_options
             or select_statement._memoized_select_entities
         ):
-            self.attributes = {}
 
             for (
                 memoized_entities
@@ -1028,9 +1067,6 @@ def create_for_statement(
                 if opt._is_compile_state:
                     opt.process_compile_state(self)
 
-        else:
-            self.attributes = {}
-
         # uncomment to print out the context.attributes structure
         # after it's been set up above
         # self._dump_option_struct()
diff --git a/test/orm/test_relationship_criteria.py b/test/orm/test_relationship_criteria.py
