ensure pool.reset event always called for reset

Added new parameter :paramref:`.PoolEvents.reset.reset_state` parameter to
the :meth:`.PoolEvents.reset` event, with deprecation logic in place that
will continue to accept event hooks using the previous set of arguments.
This indicates various state information about how the reset is taking
place and is used to allow custom reset schemes to take place with full
context given.

Within this change a fix that's also backported to 1.4 is included which
re-enables the :meth:`.PoolEvents.reset` event to continue to take place
under all circumstances, including when :class:`.Connection` has already
"reset" the connection.

The two changes together allow custom reset schemes to be implemented using
the :meth:`.PoolEvents.reset` event, instead of the
:meth:`.PoolEvents.checkin` event (which continues to function as it always
has).

Change-Id: Ie17c4f55d02beb6f570b9de6b3044baffa7d6df6
Fixes: #8717

1 files changed, 132 insertions, 1 deletions

diff --git a/test/engine/test_pool.py b/test/engine/test_pool.py
index f28da14be..9d9c3a429 100644
--- a/test/engine/test_pool.py
+++ b/test/engine/test_pool.py
@@ -9,8 +9,10 @@ from unittest.mock import patch
 import weakref
 
 import sqlalchemy as tsa
+from sqlalchemy import create_engine
 from sqlalchemy import event
 from sqlalchemy import pool
+from sqlalchemy import PoolResetState
 from sqlalchemy import select
 from sqlalchemy import testing
 from sqlalchemy.engine import default
@@ -1909,14 +1911,143 @@ class ResetOnReturnTest(PoolTestBase):
             pool.QueuePool(creator=lambda: dbapi.connect("foo.db"), **kw),
         )
 
-    def test_plain_rollback(self):
+    def _engine_fixture(self, **kw):
+        dbapi = Mock()
+
+        return dbapi, create_engine(
+            "postgresql://",
+            module=dbapi,
+            creator=lambda: dbapi.connect("foo.db"),
+            _initialize=False,
+        )
+
+    @testing.combinations("detach", "invalidate", "return")
+    def test_custom(self, extra_step):
+        dbapi, p = self._fixture(reset_on_return=None)
+
+        @event.listens_for(p, "reset")
+        def custom_reset(dbapi_conn, record, reset_state):
+            dbapi_conn.special_reset_method(reset_state)
+
+        c1 = p.connect()
+
+        if extra_step == "detach":
+            c1.detach()
+        elif extra_step == "invalidate":
+            c1.invalidate()
+        c1.close()
+
+        special_event = mock.call.special_reset_method(
+            PoolResetState(
+                transaction_was_reset=False,
+                terminate_only=extra_step == "detach",
+                asyncio_safe=True,
+            )
+        )
+        if extra_step == "detach":
+            expected = [special_event, mock.call.close()]
+        elif extra_step == "invalidate":
+            expected = [mock.call.close()]
+        else:
+            expected = [special_event]
+        eq_(dbapi.connect().mock_calls, expected)
+
+        assert not dbapi.connect().rollback.called
+        assert not dbapi.connect().commit.called
+
+    @testing.combinations(True, False, argnames="assert_w_event")
+    @testing.combinations(True, False, argnames="use_engine_transaction")
+    def test_custom_via_engine(self, assert_w_event, use_engine_transaction):
+        dbapi, engine = self._engine_fixture(reset_on_return=None)
+
+        if assert_w_event:
+
+            @event.listens_for(engine, "reset")
+            def custom_reset(dbapi_conn, record, reset_state):
+                dbapi_conn.special_reset_method(reset_state)
+
+        c1 = engine.connect()
+        if use_engine_transaction:
+            c1.begin()
+        c1.close()
+        assert dbapi.connect().rollback.called
+
+        if assert_w_event:
+            special_event = mock.call.special_reset_method(
+                PoolResetState(
+                    transaction_was_reset=use_engine_transaction,
+                    terminate_only=False,
+                    asyncio_safe=True,
+                )
+            )
+
+            if use_engine_transaction:
+                expected = [mock.call.rollback(), special_event]
+            else:
+                expected = [special_event, mock.call.rollback()]
+            eq_(dbapi.connect().mock_calls, expected)
+
+    @testing.combinations(True, False, argnames="assert_w_event")
+    def test_plain_rollback(self, assert_w_event):
         dbapi, p = self._fixture(reset_on_return="rollback")
 
+        if assert_w_event:
+
+            @event.listens_for(p, "reset")
+            def custom_reset(dbapi_conn, record, reset_state):
+                dbapi_conn.special_reset_method(reset_state)
+
         c1 = p.connect()
         c1.close()
         assert dbapi.connect().rollback.called
         assert not dbapi.connect().commit.called
 
+        if assert_w_event:
+            special_event = mock.call.special_reset_method(
+                PoolResetState(
+                    transaction_was_reset=False,
+                    terminate_only=False,
+                    asyncio_safe=True,
+                )
+            )
+
+            expected = [special_event, mock.call.rollback()]
+            eq_(dbapi.connect().mock_calls, expected)
+
+    @testing.combinations(True, False, argnames="assert_w_event")
+    @testing.combinations(True, False, argnames="use_engine_transaction")
+    def test_plain_rollback_via_engine(
+        self, assert_w_event, use_engine_transaction
+    ):
+        dbapi, engine = self._engine_fixture(reset_on_return="rollback")
+
+        if assert_w_event:
+
+            @event.listens_for(engine, "reset")
+            def custom_reset(dbapi_conn, record, reset_state):
+                dbapi_conn.special_reset_method(reset_state)
+
+        c1 = engine.connect()
+        if use_engine_transaction:
+            c1.begin()
+        c1.close()
+        assert dbapi.connect().rollback.called
+
+        if assert_w_event:
+            special_event = mock.call.special_reset_method(
+                PoolResetState(
+                    transaction_was_reset=use_engine_transaction,
+                    terminate_only=False,
+                    asyncio_safe=True,
+                )
+            )
+
+            if use_engine_transaction:
+                expected = [mock.call.rollback(), special_event]
+            else:
+                expected = [special_event, mock.call.rollback()]
+            eq_(dbapi.connect().mock_calls, expected)
+
     def test_plain_commit(self):
         dbapi, p = self._fixture(reset_on_return="commit")