restore statement substitution to before_execute()

zzzeek · zzzeek · commit f9ba830ec059 · 2021-08-20T22:17:55.000-04:00
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. The refactor in a193971 removed this feature for some reason and there were no tests in place to detect it. I don't see any indication this was planned. Fixes: sqlalchemy#6913 Change-Id: Ia77ca08aa91ab9403f19a8eb61e2a0e41aad138a
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
