Ensure synchronize_session works with lambda statements

A few places have logic that assumes the top-level statement
is the actual UPDATE or DELETE which is not the case with a
lambda.  Ensure the correct object is used.  This
fixes issues specific to both "fetch" strategy
as well as "evaluate" strategy.

Fixes: #5442
Change-Id: Ic9cc01c696c3c338d5bc79688507e6717c4c169b

1 files changed, 31 insertions, 6 deletions

diff --git a/lib/sqlalchemy/orm/persistence.py b/lib/sqlalchemy/orm/persistence.py
index cbe7bde33..45ac2442a 100644
--- a/lib/sqlalchemy/orm/persistence.py
+++ b/lib/sqlalchemy/orm/persistence.py
@@ -1888,8 +1888,18 @@ class BulkUDCompileState(CompileState):
                 from_=err,
             )
 
-        if statement.__visit_name__ == "update":
-            resolved_values = cls._get_resolved_values(mapper, statement)
+        if statement.__visit_name__ == "lambda_element":
+            # ._resolved is called on every LambdaElement in order to
+            # generate the cache key, so this access does not add
+            # additional expense
+            effective_statement = statement._resolved
+        else:
+            effective_statement = statement
+
+        if effective_statement.__visit_name__ == "update":
+            resolved_values = cls._get_resolved_values(
+                mapper, effective_statement
+            )
             value_evaluators = {}
             resolved_keys_as_propnames = cls._resolved_keys_as_propnames(
                 mapper, resolved_values
@@ -2012,10 +2022,20 @@ class BulkUDCompileState(CompileState):
 
         value_evaluators = _EMPTY_DICT
 
-        if statement.__visit_name__ == "update":
+        if statement.__visit_name__ == "lambda_element":
+            # ._resolved is called on every LambdaElement in order to
+            # generate the cache key, so this access does not add
+            # additional expense
+            effective_statement = statement._resolved
+        else:
+            effective_statement = statement
+
+        if effective_statement.__visit_name__ == "update":
             target_cls = mapper.class_
             evaluator_compiler = evaluator.EvaluatorCompiler(target_cls)
-            resolved_values = cls._get_resolved_values(mapper, statement)
+            resolved_values = cls._get_resolved_values(
+                mapper, effective_statement
+            )
             resolved_keys_as_propnames = cls._resolved_keys_as_propnames(
                 mapper, resolved_values
             )
@@ -2073,8 +2093,12 @@ class BulkORMUpdate(UpdateDMLState, BulkUDCompileState):
         elif statement._values:
             new_stmt._values = self._resolved_values
 
+        # if we are against a lambda statement we might not be the
+        # topmost object that received per-execute annotations
+        top_level_stmt = compiler.statement
         if (
-            statement._annotations.get("synchronize_session", None) == "fetch"
+            top_level_stmt._annotations.get("synchronize_session", None)
+            == "fetch"
             and compiler.dialect.full_returning
         ):
             new_stmt = new_stmt.returning(*mapper.primary_key)
@@ -2187,9 +2211,10 @@ class BulkORMDelete(DeleteDMLState, BulkUDCompileState):
             "parentmapper", None
         )
 
+        top_level_stmt = compiler.statement
         if (
             mapper
-            and statement._annotations.get("synchronize_session", None)
+            and top_level_stmt._annotations.get("synchronize_session", None)
             == "fetch"
             and compiler.dialect.full_returning
         ):