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

3 files changed, 169 insertions, 53 deletions

diff --git a/lib/sqlalchemy/pool/__init__.py b/lib/sqlalchemy/pool/__init__.py
index c86d3dded..8eb2ae4ab 100644
--- a/lib/sqlalchemy/pool/__init__.py
+++ b/lib/sqlalchemy/pool/__init__.py
@@ -29,6 +29,7 @@ from .base import ConnectionPoolEntry as ConnectionPoolEntry
 from .base import ManagesConnection as ManagesConnection
 from .base import Pool as Pool
 from .base import PoolProxiedConnection as PoolProxiedConnection
+from .base import PoolResetState as PoolResetState
 from .base import reset_commit as reset_commit
 from .base import reset_none as reset_none
 from .base import reset_rollback as reset_rollback
diff --git a/lib/sqlalchemy/pool/base.py b/lib/sqlalchemy/pool/base.py
index 41f2f03b2..18ff66e5d 100644
--- a/lib/sqlalchemy/pool/base.py
+++ b/lib/sqlalchemy/pool/base.py
@@ -13,6 +13,7 @@
 from __future__ import annotations
 
 from collections import deque
+import dataclasses
 from enum import Enum
 import threading
 import time
@@ -46,6 +47,51 @@ if TYPE_CHECKING:
     from ..sql._typing import _InfoType
 
 
+@dataclasses.dataclass(frozen=True)
+class PoolResetState:
+    """describes the state of a DBAPI connection as it is being passed to
+    the :meth:`.PoolEvents.reset` connection pool event.
+
+    .. versionadded:: 2.0.0b3
+
+    """
+
+    __slots__ = ("transaction_was_reset", "terminate_only", "asyncio_safe")
+
+    transaction_was_reset: bool
+    """Indicates if the transaction on the DBAPI connection was already
+    essentially "reset" back by the :class:`.Connection` object.
+
+    This boolean is True if the :class:`.Connection` had transactional
+    state present upon it, which was then not closed using the
+    :meth:`.Connection.rollback` or :meth:`.Connection.commit` method;
+    instead, the transaction was closed inline within the
+    :meth:`.Connection.close` method so is guaranteed to remain non-present
+    when this event is reached.
+
+    """
+
+    terminate_only: bool
+    """indicates if the connection is to be immediately terminated and
+    not checked in to the pool.
+
+    This occurs for connections that were invalidated, as well as asyncio
+    connections that were not cleanly handled by the calling code that
+    are instead being garbage collected.   In the latter case,
+    operations can't be safely run on asyncio connections within garbage
+    collection as there is not necessarily an event loop present.
+
+    """
+
+    asyncio_safe: bool
+    """Indicates if the reset operation is occurring within a scope where
+    an enclosing event loop is expected to be present for asyncio applications.
+
+    Will be False in the case that the connection is being garbage collected.
+
+    """
+
+
 class ResetStyle(Enum):
     """Describe options for "reset on return" behaviors."""
 
@@ -168,39 +214,46 @@ class Pool(log.Identified, event.EventTarget):
              logging.
 
         :param reset_on_return: Determine steps to take on
-          connections as they are returned to the pool, which were
-          not otherwise handled by a :class:`_engine.Connection`.
-
-          reset_on_return can have any of these values:
-
-          * ``"rollback"`` - call rollback() on the connection,
-            to release locks and transaction resources.
-            This is the default value.  The vast majority
-            of use cases should leave this value set.
-          * ``True`` - same as 'rollback', this is here for
-            backwards compatibility.
-          * ``"commit"`` - call commit() on the connection,
-            to release locks and transaction resources.
-            A commit here may be desirable for databases that
-            cache query plans if a commit is emitted,
-            such as Microsoft SQL Server.  However, this
-            value is more dangerous than 'rollback' because
-            any data changes present on the transaction
-            are committed unconditionally.
-          * ``None`` - don't do anything on the connection.
-            This setting may be appropriate if the database / DBAPI
-            works in pure "autocommit" mode at all times, or if the
-            application uses the :class:`_engine.Engine` with consistent
-            connectivity patterns.   See the section
-            :ref:`pool_reset_on_return` for more details.
-
-          * ``False`` - same as None, this is here for
-            backwards compatibility.
+         connections as they are returned to the pool, which were
+         not otherwise handled by a :class:`_engine.Connection`.
+         Available from :func:`_sa.create_engine` via the
+         :paramref:`_sa.create_engine.pool_reset_on_return` parameter.
+
+         :paramref:`_pool.Pool.reset_on_return` can have any of these values:
+
+         * ``"rollback"`` - call rollback() on the connection,
+           to release locks and transaction resources.
+           This is the default value.  The vast majority
+           of use cases should leave this value set.
+         * ``"commit"`` - call commit() on the connection,
+           to release locks and transaction resources.
+           A commit here may be desirable for databases that
+           cache query plans if a commit is emitted,
+           such as Microsoft SQL Server.  However, this
+           value is more dangerous than 'rollback' because
+           any data changes present on the transaction
+           are committed unconditionally.
+         * ``None`` - don't do anything on the connection.
+           This setting may be appropriate if the database / DBAPI
+           works in pure "autocommit" mode at all times, or if
+           a custom reset handler is established using the
+           :meth:`.PoolEvents.reset` event handler.
+
+         * ``True`` - same as 'rollback', this is here for
+           backwards compatibility.
+         * ``False`` - same as None, this is here for
+           backwards compatibility.
+
+         For further customization of reset on return, the
+         :meth:`.PoolEvents.reset` event hook may be used which can perform
+         any connection activity desired on reset.
 
          .. seealso::
 
             :ref:`pool_reset_on_return`
 
+            :meth:`.PoolEvents.reset`
+
         :param events: a list of 2-tuples, each of the form
          ``(callable, target)`` which will be passed to :func:`.event.listen`
          upon construction.   Provided here so that event listeners
@@ -671,7 +724,9 @@ class _ConnectionRecord(ConnectionPoolEntry):
 
         rec.fairy_ref = ref = weakref.ref(
             fairy,
-            lambda ref: _finalize_fairy(None, rec, pool, ref, echo, True)
+            lambda ref: _finalize_fairy(
+                None, rec, pool, ref, echo, transaction_was_reset=False
+            )
             if _finalize_fairy
             else None,
         )
@@ -864,7 +919,7 @@ def _finalize_fairy(
         weakref.ref[_ConnectionFairy]
     ],  # this is None when called directly, not by the gc
     echo: Optional[log._EchoFlagType],
-    reset: bool = True,
+    transaction_was_reset: bool = False,
     fairy: Optional[_ConnectionFairy] = None,
 ) -> None:
     """Cleanup for a :class:`._ConnectionFairy` whether or not it's already
@@ -906,11 +961,7 @@ def _finalize_fairy(
     if dbapi_connection is not None:
         if connection_record and echo:
             pool.logger.debug(
-                "Connection %r being returned to pool%s",
-                dbapi_connection,
-                ", transaction state was already reset by caller"
-                if not reset
-                else "",
+                "Connection %r being returned to pool", dbapi_connection
             )
 
         try:
@@ -923,8 +974,13 @@ def _finalize_fairy(
                     echo,
                 )
             assert fairy.dbapi_connection is dbapi_connection
-            if reset and can_manipulate_connection:
-                fairy._reset(pool)
+
+            fairy._reset(
+                pool,
+                transaction_was_reset=transaction_was_reset,
+                terminate_only=detach,
+                asyncio_safe=can_manipulate_connection,
+            )
 
             if detach:
                 if connection_record:
@@ -1293,29 +1349,55 @@ class _ConnectionFairy(PoolProxiedConnection):
     def _checkout_existing(self) -> _ConnectionFairy:
         return _ConnectionFairy._checkout(self._pool, fairy=self)
 
-    def _checkin(self, reset: bool = True) -> None:
+    def _checkin(self, transaction_was_reset: bool = False) -> None:
         _finalize_fairy(
             self.dbapi_connection,
             self._connection_record,
             self._pool,
             None,
             self._echo,
-            reset=reset,
+            transaction_was_reset=transaction_was_reset,
             fairy=self,
         )
 
     def _close(self) -> None:
         self._checkin()
 
-    def _reset(self, pool: Pool) -> None:
+    def _reset(
+        self,
+        pool: Pool,
+        transaction_was_reset: bool,
+        terminate_only: bool,
+        asyncio_safe: bool,
+    ) -> None:
         if pool.dispatch.reset:
-            pool.dispatch.reset(self.dbapi_connection, self._connection_record)
+            pool.dispatch.reset(
+                self.dbapi_connection,
+                self._connection_record,
+                PoolResetState(
+                    transaction_was_reset=transaction_was_reset,
+                    terminate_only=terminate_only,
+                    asyncio_safe=asyncio_safe,
+                ),
+            )
+
+        if not asyncio_safe:
+            return
+
         if pool._reset_on_return is reset_rollback:
-            if self._echo:
-                pool.logger.debug(
-                    "Connection %s rollback-on-return", self.dbapi_connection
-                )
-            pool._dialect.do_rollback(self)
+            if transaction_was_reset:
+                if self._echo:
+                    pool.logger.debug(
+                        "Connection %s reset, transaction already reset",
+                        self.dbapi_connection,
+                    )
+            else:
+                if self._echo:
+                    pool.logger.debug(
+                        "Connection %s rollback-on-return",
+                        self.dbapi_connection,
+                    )
+                pool._dialect.do_rollback(self)
         elif pool._reset_on_return is reset_commit:
             if self._echo:
                 pool.logger.debug(
@@ -1396,7 +1478,7 @@ class _ConnectionFairy(PoolProxiedConnection):
         if self._counter == 0:
             self._checkin()
 
-    def _close_no_reset(self) -> None:
+    def _close_special(self, transaction_reset: bool = False) -> None:
         self._counter -= 1
         if self._counter == 0:
-            self._checkin(reset=False)
+            self._checkin(transaction_was_reset=transaction_reset)
diff --git a/lib/sqlalchemy/pool/events.py b/lib/sqlalchemy/pool/events.py
index 47ab106d7..0bf69420e 100644
--- a/lib/sqlalchemy/pool/events.py
+++ b/lib/sqlalchemy/pool/events.py
@@ -15,6 +15,7 @@ from typing import Union
 from .base import ConnectionPoolEntry
 from .base import Pool
 from .base import PoolProxiedConnection
+from .base import PoolResetState
 from .. import event
 from .. import util
 
@@ -189,22 +190,46 @@ class PoolEvents(event.Events[Pool]):
 
         """
 
+    @event._legacy_signature(
+        "2.0",
+        ["dbapi_connection", "connection_record"],
+        lambda dbapi_connection, connection_record, reset_state: (
+            dbapi_connection,
+            connection_record,
+        ),
+    )
     def reset(
         self,
         dbapi_connection: DBAPIConnection,
         connection_record: ConnectionPoolEntry,
+        reset_state: PoolResetState,
     ) -> None:
         """Called before the "reset" action occurs for a pooled connection.
 
         This event represents
         when the ``rollback()`` method is called on the DBAPI connection
-        before it is returned to the pool.  The behavior of "reset" can
-        be controlled, including disabled, using the ``reset_on_return``
-        pool argument.
-
+        before it is returned to the pool or discarded.
+        A custom "reset" strategy may be implemented using this event hook,
+        which may also be combined with disabling the default "reset"
+        behavior using the :paramref:`_pool.Pool.reset_on_return` parameter.
+
+        The primary difference between the :meth:`_events.PoolEvents.reset` and
+        :meth:`_events.PoolEvents.checkin` events are that
+        :meth:`_events.PoolEvents.reset` is called not just for pooled
+        connections that are being returned to the pool, but also for
+        connections that were detached using the
+        :meth:`_engine.Connection.detach` method as well as asyncio connections
+        that are being discarded due to garbage collection taking place on
+        connections before the connection was checked in.
+
+        Note that the event **is not** invoked for connections that were
+        invalidated using :meth:`_engine.Connection.invalidate`.    These
+        events may be intercepted using the :meth:`.PoolEvents.soft_invalidate`
+        and :meth:`.PoolEvents.invalidate` event hooks, and all "connection
+        close" events may be intercepted using :meth:`.PoolEvents.close`.
 
         The :meth:`_events.PoolEvents.reset` event is usually followed by the
-        :meth:`_events.PoolEvents.checkin` event is called, except in those
+        :meth:`_events.PoolEvents.checkin` event, except in those
         cases where the connection is discarded immediately after reset.
 
         :param dbapi_connection: a DBAPI connection.
@@ -213,8 +238,16 @@ class PoolEvents(event.Events[Pool]):
         :param connection_record: the :class:`.ConnectionPoolEntry` managing
          the DBAPI connection.
 
+        :param reset_state: :class:`.PoolResetState` instance which provides
+         information about the circumstances under which the connection
+         is being reset.
+
+         .. versionadded:: 2.0
+
         .. seealso::
 
+            :ref:`pool_reset_on_return`
+
             :meth:`_events.ConnectionEvents.rollback`
 
             :meth:`_events.ConnectionEvents.commit`