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

10 files changed, 306 insertions, 59 deletions

diff --git a/lib/sqlalchemy/__init__.py b/lib/sqlalchemy/__init__.py
index 5823239dc..b22d7ca8c 100644
--- a/lib/sqlalchemy/__init__.py
+++ b/lib/sqlalchemy/__init__.py
@@ -52,6 +52,7 @@ from .pool import (
 from .pool import NullPool as NullPool
 from .pool import Pool as Pool
 from .pool import PoolProxiedConnection as PoolProxiedConnection
+from .pool import PoolResetState as PoolResetState
 from .pool import QueuePool as QueuePool
 from .pool import SingletonThreadPool as SingleonThreadPool
 from .pool import StaticPool as StaticPool
diff --git a/lib/sqlalchemy/dialects/mssql/base.py b/lib/sqlalchemy/dialects/mssql/base.py
index 5d42d98e3..a338ba27a 100644
--- a/lib/sqlalchemy/dialects/mssql/base.py
+++ b/lib/sqlalchemy/dialects/mssql/base.py
@@ -480,6 +480,61 @@ different isolation level settings.  See the discussion at
 
     :ref:`dbapi_autocommit`
 
+.. _mssql_reset_on_return:
+
+Temporary Table / Resource Reset for Connection Pooling
+-------------------------------------------------------
+
+The :class:`.QueuePool` connection pool implementation used
+by the SQLAlchemy :class:`.Engine` object includes
+:ref:`reset on return <pool_reset_on_return>` behavior that will invoke
+the DBAPI ``.rollback()`` method when connections are returned to the pool.
+While this rollback will clear out the immediate state used by the previous
+transaction, it does not cover a wider range of session-level state, including
+temporary tables as well as other server state such as prepared statement
+handles and statement caches.   An undocumented SQL Server procedure known
+as ``sp_reset_connection`` is known to be a workaround for this issue which
+will reset most of the session state that builds up on a connection, including
+temporary tables.
+
+To install ``sp_reset_connection`` as the means of performing reset-on-return,
+the :meth:`.PoolEvents.reset` event hook may be used, as demonstrated in the
+example below. The :paramref:`_sa.create_engine.pool_reset_on_return` parameter
+is set to ``None`` so that the custom scheme can replace the default behavior
+completely.   The custom hook implementation calls ``.rollback()`` in any case,
+as it's usually important that the DBAPI's own tracking of commit/rollback
+will remain consistent with the state of the transaction::
+
+    from sqlalchemy import create_engine
+    from sqlalchemy import event
+
+    mssql_engine = create_engine(
+        "mssql+pyodbc://scott:tiger^5HHH@mssql2017:1433/test?driver=ODBC+Driver+17+for+SQL+Server",
+
+        # disable default reset-on-return scheme
+        pool_reset_on_return=None,
+    )
+
+
+    @event.listens_for(mssql_engine, "reset")
+    def _reset_mssql(dbapi_connection, connection_record, reset_state):
+        if not reset_state.terminate_only:
+            dbapi_connection.execute("{call sys.sp_reset_connection}")
+
+        # so that the DBAPI itself knows that the connection has been
+        # reset
+        dbapi_connection.rollback()
+
+.. versionchanged:: 2.0.0b3  Added additional state arguments to
+   the :meth:`.PoolEvents.reset` event and additionally ensured the event
+   is invoked for all "reset" occurrences, so that it's appropriate
+   as a place for custom "reset" handlers.   Previous schemes which
+   use the :meth:`.PoolEvents.checkin` handler remain usable as well.
+
+.. seealso::
+
+    :ref:`pool_reset_on_return` - in the :ref:`pooling_toplevel` documentation
+
 Nullability
 -----------
 MSSQL has support for three levels of column nullability. The default
diff --git a/lib/sqlalchemy/dialects/postgresql/base.py b/lib/sqlalchemy/dialects/postgresql/base.py
index 5eab8f47c..482e36594 100644
--- a/lib/sqlalchemy/dialects/postgresql/base.py
+++ b/lib/sqlalchemy/dialects/postgresql/base.py
@@ -233,6 +233,70 @@ SERIALIZABLE isolation.
 .. versionadded:: 1.4 added support for the ``postgresql_readonly``
    and ``postgresql_deferrable`` execution options.
 
+.. _postgresql_reset_on_return:
+
+Temporary Table / Resource Reset for Connection Pooling
+-------------------------------------------------------
+
+The :class:`.QueuePool` connection pool implementation used
+by the SQLAlchemy :class:`.Engine` object includes
+:ref:`reset on return <pool_reset_on_return>` behavior that will invoke
+the DBAPI ``.rollback()`` method when connections are returned to the pool.
+While this rollback will clear out the immediate state used by the previous
+transaction, it does not cover a wider range of session-level state, including
+temporary tables as well as other server state such as prepared statement
+handles and statement caches.   The PostgreSQL database includes a variety
+of commands which may be used to reset this state, including
+``DISCARD``, ``RESET``, ``DEALLOCATE``, and ``UNLISTEN``.
+
+
+To install
+one or more of these commands as the means of performing reset-on-return,
+the :meth:`.PoolEvents.reset` event hook may be used, as demonstrated
+in the example below. The implementation
+will end transactions in progress as well as discard temporary tables
+using the ``CLOSE``, ``RESET`` and ``DISCARD`` commands; see the PostgreSQL
+documentation for background on what each of these statements do.
+
+The :paramref:`_sa.create_engine.pool_reset_on_return` parameter
+is set to ``None`` so that the custom scheme can replace the default behavior
+completely.   The custom hook implementation calls ``.rollback()`` in any case,
+as it's usually important that the DBAPI's own tracking of commit/rollback
+will remain consistent with the state of the transaction::
+
+
+    from sqlalchemy import create_engine
+    from sqlalchemy import event
+
+    postgresql_engine = create_engine(
+        "postgresql+pyscopg2://scott:tiger@hostname/dbname",
+
+        # disable default reset-on-return scheme
+        pool_reset_on_return=None,
+    )
+
+
+    @event.listens_for(postgresql_engine, "reset")
+    def _reset_mssql(dbapi_connection, connection_record, reset_state):
+        if not reset_state.terminate_only:
+            dbapi_connection.execute("CLOSE ALL")
+            dbapi_connection.execute("RESET ALL")
+            dbapi_connection.execute("DISCARD TEMP")
+
+        # so that the DBAPI itself knows that the connection has been
+        # reset
+        dbapi_connection.rollback()
+
+.. versionchanged:: 2.0.0b3  Added additional state arguments to
+   the :meth:`.PoolEvents.reset` event and additionally ensured the event
+   is invoked for all "reset" occurrences, so that it's appropriate
+   as a place for custom "reset" handlers.   Previous schemes which
+   use the :meth:`.PoolEvents.checkin` handler remain usable as well.
+
+.. seealso::
+
+    :ref:`pool_reset_on_return` - in the :ref:`pooling_toplevel` documentation
+
 .. _postgresql_alternate_search_path:
 
 Setting Alternate Search Paths on Connect
diff --git a/lib/sqlalchemy/engine/base.py b/lib/sqlalchemy/engine/base.py
index 1ec307297..8cbc0163c 100644
--- a/lib/sqlalchemy/engine/base.py
+++ b/lib/sqlalchemy/engine/base.py
@@ -1220,7 +1220,9 @@ class Connection(ConnectionEventsTarget, inspection.Inspectable["Inspector"]):
             # as we just closed the transaction, close the connection
             # pool connection without doing an additional reset
             if skip_reset:
-                cast("_ConnectionFairy", conn)._close_no_reset()
+                cast("_ConnectionFairy", conn)._close_special(
+                    transaction_reset=True
+                )
             else:
                 conn.close()
 
diff --git a/lib/sqlalchemy/engine/create.py b/lib/sqlalchemy/engine/create.py
index a9b388d71..ef5463972 100644
--- a/lib/sqlalchemy/engine/create.py
+++ b/lib/sqlalchemy/engine/create.py
@@ -469,7 +469,7 @@ def create_engine(url: Union[str, "_url.URL"], **kwargs: Any) -> Engine:
 
         .. seealso::
 
-            :paramref:`_pool.Pool.reset_on_return`
+            :ref:`pool_reset_on_return`
 
     :param pool_timeout=30: number of seconds to wait before giving
         up on getting a connection from the pool. This is only used
diff --git a/lib/sqlalchemy/event/legacy.py b/lib/sqlalchemy/event/legacy.py
index 8ae9b7035..3d4341410 100644
--- a/lib/sqlalchemy/event/legacy.py
+++ b/lib/sqlalchemy/event/legacy.py
@@ -194,9 +194,9 @@ def _version_signature_changes(
 ) -> str:
     since, args, conv = dispatch_collection.legacy_signatures[0]
     return (
-        "\n.. deprecated:: %(since)s\n"
-        "    The :class:`.%(clsname)s.%(event_name)s` event now accepts the \n"
-        "    arguments ``%(named_event_arguments)s%(has_kw_arguments)s``.\n"
+        "\n.. versionchanged:: %(since)s\n"
+        "    The :meth:`.%(clsname)s.%(event_name)s` event now accepts the \n"
+        "    arguments %(named_event_arguments)s%(has_kw_arguments)s.\n"
         "    Support for listener functions which accept the previous \n"
         '    argument signature(s) listed above as "deprecated" will be \n'
         "    removed in a future release."
@@ -204,7 +204,15 @@ def _version_signature_changes(
             "since": since,
             "clsname": parent_dispatch_cls.__name__,
             "event_name": dispatch_collection.name,
-            "named_event_arguments": ", ".join(dispatch_collection.arg_names),
+            "named_event_arguments": ", ".join(
+                ":paramref:`.%(clsname)s.%(event_name)s.%(param_name)s`"
+                % {
+                    "clsname": parent_dispatch_cls.__name__,
+                    "event_name": dispatch_collection.name,
+                    "param_name": param_name,
+                }
+                for param_name in dispatch_collection.arg_names
+            ),
             "has_kw_arguments": ", **kw" if dispatch_collection.has_kw else "",
         }
     )
diff --git a/lib/sqlalchemy/events.py b/lib/sqlalchemy/events.py
index 4a8a52337..4e81ceec5 100644
--- a/lib/sqlalchemy/events.py
+++ b/lib/sqlalchemy/events.py
@@ -11,6 +11,7 @@ from __future__ import annotations
 
 from .engine.events import ConnectionEvents
 from .engine.events import DialectEvents
+from .pool import PoolResetState
 from .pool.events import PoolEvents
 from .sql.base import SchemaEventTarget
 from .sql.events import DDLEvents
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`