Merge "restore statement substitution to before_execute()"

zzzeek · Gerrit Code Review · commit b5822c2d4a10 · 2021-08-23T16:37:10.000Z
diff --git a/doc/build/changelog/unreleased_14/6913.rst b/doc/build/changelog/unreleased_14/6913.rst
@@ -0,0 +1,9 @@
+.. change::
+    :tags: bug, engine, regression
+    :tickets: 6913
+
+    Fixed issue where the ability of the
+    :meth:`_engine.ConnectionEvents.before_execute` method to alter the SQL
+    statement object passed, returning the new object to be invoked, was
+    inadvertently removed. This behavior has been restored.
+
diff --git a/lib/sqlalchemy/engine/base.py b/lib/sqlalchemy/engine/base.py
@@ -1287,6 +1287,7 @@ def _execute_default(
 
         if self._has_events or self.engine._has_events:
             (
+                default,
                 distilled_params,
                 event_multiparams,
                 event_params,
@@ -1335,6 +1336,7 @@ def _execute_ddl(self, ddl, multiparams, params, execution_options):
 
         if self._has_events or self.engine._has_events:
             (
+                ddl,
                 distilled_params,
                 event_multiparams,
                 event_params,
@@ -1399,7 +1401,7 @@ def _invoke_before_exec_event(
         else:
             distilled_params = []
 
-        return distilled_params, event_multiparams, event_params
+        return elem, distilled_params, event_multiparams, event_params
 
     def _execute_clauseelement(
         self, elem, multiparams, params, execution_options
@@ -1415,6 +1417,7 @@ def _execute_clauseelement(
         has_events = self._has_events or self.engine._has_events
         if has_events:
             (
+                elem,
                 distilled_params,
                 event_multiparams,
                 event_params,
@@ -1492,6 +1495,7 @@ def _execute_compiled(
 
         if self._has_events or self.engine._has_events:
             (
+                compiled,
                 distilled_params,
                 event_multiparams,
                 event_params,
@@ -1536,6 +1540,7 @@ def _exec_driver_sql(
         if not future:
             if self._has_events or self.engine._has_events:
                 (
+                    statement,
                     distilled_params,
                     event_multiparams,
                     event_params,
diff --git a/test/engine/test_deprecations.py b/test/engine/test_deprecations.py
@@ -1687,6 +1687,17 @@ def after_execute(
                 )
             eq_(result.all(), [("15",)])
 
+    @testing.only_on("sqlite")
+    def test_modify_statement_string(self, connection):
+        @event.listens_for(connection, "before_execute", retval=True)
+        def _modify(
+            conn, clauseelement, multiparams, params, execution_options
+        ):
+            return clauseelement.replace("hi", "there"), multiparams, params
+
+        with _string_deprecation_expect():
+            eq_(connection.scalar("select 'hi'"), "there")
+
     def test_retval_flag(self):
         canary = []
 
diff --git a/test/engine/test_execute.py b/test/engine/test_execute.py
@@ -31,6 +31,7 @@
 from sqlalchemy.pool import QueuePool
 from sqlalchemy.sql import column
 from sqlalchemy.sql import literal
+from sqlalchemy.sql.elements import literal_column
 from sqlalchemy.testing import assert_raises
 from sqlalchemy.testing import assert_raises_message
 from sqlalchemy.testing import config
@@ -1771,6 +1772,56 @@ def before_execute(
             with e1.connect() as conn:
                 conn.execute(select(literal("1")))
 
+    @testing.only_on("sqlite")
+    def test_dont_modify_statement_driversql(self, connection):
+        m1 = mock.Mock()
+
+        @event.listens_for(connection, "before_execute", retval=True)
+        def _modify(
+            conn, clauseelement, multiparams, params, execution_options
+        ):
+            m1.run_event()
+            return clauseelement.replace("hi", "there"), multiparams, params
+
+        # the event does not take effect for the "driver SQL" option
+        eq_(connection.exec_driver_sql("select 'hi'").scalar(), "hi")
+
+        # event is not called at all
+        eq_(m1.mock_calls, [])
+
+    @testing.combinations((True,), (False,), argnames="future")
+    @testing.only_on("sqlite")
+    def test_modify_statement_internal_driversql(self, connection, future):
+        m1 = mock.Mock()
+
+        @event.listens_for(connection, "before_execute", retval=True)
+        def _modify(
+            conn, clauseelement, multiparams, params, execution_options
+        ):
+            m1.run_event()
+            return clauseelement.replace("hi", "there"), multiparams, params
+
+        eq_(
+            connection._exec_driver_sql(
+                "select 'hi'", [], {}, {}, future=future
+            ).scalar(),
+            "hi" if future else "there",
+        )
+
+        if future:
+            eq_(m1.mock_calls, [])
+        else:
+            eq_(m1.mock_calls, [call.run_event()])
+
+    def test_modify_statement_clauseelement(self, connection):
+        @event.listens_for(connection, "before_execute", retval=True)
+        def _modify(
+            conn, clauseelement, multiparams, params, execution_options
+        ):
+            return select(literal_column("'there'")), multiparams, params
+
+        eq_(connection.scalar(select(literal_column("'hi'"))), "there")
+
     def test_argument_format_execute(self, testing_engine):
         def before_execute(
             conn, clauseelement, multiparams, params, execution_options
