From b96321ae79a0366c33ca739e6e67aaf5f4420db4 Mon Sep 17 00:00:00 2001
From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Tue, 25 Oct 2022 09:10:09 -0400
Subject: Support result.close() for all iterator patterns

This change contains new features for 2.0 only as well as some
behaviors that will be backported to 1.4.

For 1.4 and 2.0:

Fixed issue where the underlying DBAPI cursor would not be closed when
using :class:`_orm.Query` with :meth:`_orm.Query.yield_per` and direct
iteration, if a user-defined exception case were raised within the
iteration process, interrupting the iterator. This would lead to the usual
MySQL-related issues with server side cursors out of sync.

For 1.4 only:

A similar scenario can occur when using :term:`2.x` executions with direct
use of :class:`.Result`, in that case the end-user code has access to the
:class:`.Result` itself and should call :meth:`.Result.close` directly.
Version 2.0 will feature context-manager calling patterns to address this
use case.  However within the 1.4 scope, ensured that ``.close()`` methods
are available on all :class:`.Result` implementations including
:class:`.ScalarResult`, :class:`.MappingResult`.

For 2.0 only:

To better support the use case of iterating :class:`.Result` and
:class:`.AsyncResult` objects where user-defined exceptions may interrupt
the iteration, both objects as well as variants such as
:class:`.ScalarResult`, :class:`.MappingResult`,
:class:`.AsyncScalarResult`, :class:`.AsyncMappingResult` now support
context manager usage, where the result will be closed at the end of
iteration.

Corrected various typing issues within the engine and async engine
packages.

Fixes: #8710
Change-Id: I3166328bfd3900957eb33cbf1061d0495c9df670
---
 lib/sqlalchemy/engine/result.py | 59 +++++++++++++++++++++++++++++++++++++++--
 1 file changed, 57 insertions(+), 2 deletions(-)

(limited to 'lib/sqlalchemy/engine')

diff --git a/lib/sqlalchemy/engine/result.py b/lib/sqlalchemy/engine/result.py
index 05ca17063..bcd2f0ea9 100644
--- a/lib/sqlalchemy/engine/result.py
+++ b/lib/sqlalchemy/engine/result.py
@@ -928,6 +928,12 @@ class Result(_WithKeys, ResultInternal[Row[_TP]]):
     def __init__(self, cursor_metadata: ResultMetaData):
         self._metadata = cursor_metadata
 
+    def __enter__(self) -> Result[_TP]:
+        return self
+
+    def __exit__(self, type_: Any, value: Any, traceback: Any) -> None:
+        self.close()
+
     def close(self) -> None:
         """close this :class:`_result.Result`.
 
@@ -950,6 +956,19 @@ class Result(_WithKeys, ResultInternal[Row[_TP]]):
         """
         self._soft_close(hard=True)
 
+    @property
+    def _soft_closed(self) -> bool:
+        raise NotImplementedError()
+
+    @property
+    def closed(self) -> bool:
+        """return True if this :class:`.Result` reports .closed
+
+        .. versionadded:: 1.4.43
+
+        """
+        raise NotImplementedError()
+
     @_generative
     def yield_per(self: SelfResult, num: int) -> SelfResult:
         """Configure the row-fetching strategy to fetch ``num`` rows at a time.
@@ -1574,6 +1593,12 @@ class FilterResult(ResultInternal[_R]):
 
     _real_result: Result[Any]
 
+    def __enter__(self: SelfFilterResult) -> SelfFilterResult:
+        return self
+
+    def __exit__(self, type_: Any, value: Any, traceback: Any) -> None:
+        self._real_result.__exit__(type_, value, traceback)
+
     @_generative
     def yield_per(self: SelfFilterResult, num: int) -> SelfFilterResult:
         """Configure the row-fetching strategy to fetch ``num`` rows at a time.
@@ -1599,6 +1624,27 @@ class FilterResult(ResultInternal[_R]):
     def _soft_close(self, hard: bool = False) -> None:
         self._real_result._soft_close(hard=hard)
 
+    @property
+    def _soft_closed(self) -> bool:
+        return self._real_result._soft_closed
+
+    @property
+    def closed(self) -> bool:
+        """return True if the underlying :class:`.Result` reports .closed
+
+        .. versionadded:: 1.4.43
+
+        """
+        return self._real_result.closed
+
+    def close(self) -> None:
+        """Close this :class:`.FilterResult`.
+
+        .. versionadded:: 1.4.43
+
+        """
+        self._real_result.close()
+
     @property
     def _attributes(self) -> Dict[Any, Any]:
         return self._real_result._attributes
@@ -2172,7 +2218,7 @@ class IteratorResult(Result[_TP]):
         self,
         cursor_metadata: ResultMetaData,
         iterator: Iterator[_InterimSupportsScalarsRowType],
-        raw: Optional[Any] = None,
+        raw: Optional[Result[Any]] = None,
         _source_supports_scalars: bool = False,
     ):
         self._metadata = cursor_metadata
@@ -2180,6 +2226,15 @@ class IteratorResult(Result[_TP]):
         self.raw = raw
         self._source_supports_scalars = _source_supports_scalars
 
+    @property
+    def closed(self) -> bool:
+        """return True if this :class:`.IteratorResult` has been closed
+
+        .. versionadded:: 1.4.43
+
+        """
+        return self._hard_closed
+
     def _soft_close(self, hard: bool = False, **kw: Any) -> None:
         if hard:
             self._hard_closed = True
@@ -2262,7 +2317,7 @@ class ChunkedIteratorResult(IteratorResult[_TP]):
             [Optional[int]], Iterator[Sequence[_InterimRowType[_R]]]
         ],
         source_supports_scalars: bool = False,
-        raw: Optional[Any] = None,
+        raw: Optional[Result[Any]] = None,
         dynamic_yield_per: bool = False,
     ):
         self._metadata = cursor_metadata
-- 
cgit v1.2.1