Merge "ensure pool.reset event always called for reset" into main

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`