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

1 files changed, 125 insertions, 0 deletions

diff --git a/test/sql/test_resultset.py b/test/sql/test_resultset.py
index 4f776e300..fa86d75ee 100644
--- a/test/sql/test_resultset.py
+++ b/test/sql/test_resultset.py
@@ -53,6 +53,7 @@ from sqlalchemy.testing import expect_raises_message
 from sqlalchemy.testing import fixtures
 from sqlalchemy.testing import in_
 from sqlalchemy.testing import is_
+from sqlalchemy.testing import is_false
 from sqlalchemy.testing import is_true
 from sqlalchemy.testing import le_
 from sqlalchemy.testing import mock
@@ -2033,6 +2034,89 @@ class CursorResultTest(fixtures.TablesTest):
             partition = next(result.partitions())
             eq_(len(partition), value)
 
+    @testing.fixture
+    def autoclose_row_fixture(self, connection):
+        users = self.tables.users
+        connection.execute(
+            users.insert(),
+            [
+                {"user_id": 1, "name": "u1"},
+                {"user_id": 2, "name": "u2"},
+                {"user_id": 3, "name": "u3"},
+                {"user_id": 4, "name": "u4"},
+                {"user_id": 5, "name": "u5"},
+            ],
+        )
+
+    @testing.fixture(params=["plain", "scalars", "mapping"])
+    def result_fixture(self, request, connection):
+        users = self.tables.users
+
+        result_type = request.param
+
+        if result_type == "plain":
+            result = connection.execute(select(users))
+        elif result_type == "scalars":
+            result = connection.scalars(select(users))
+        elif result_type == "mapping":
+            result = connection.execute(select(users)).mappings()
+        else:
+            assert False
+
+        return result
+
+    def test_results_can_close(self, autoclose_row_fixture, result_fixture):
+        """test #8710"""
+
+        r1 = result_fixture
+
+        is_false(r1.closed)
+        is_false(r1._soft_closed)
+
+        r1._soft_close()
+        is_false(r1.closed)
+        is_true(r1._soft_closed)
+
+        r1.close()
+        is_true(r1.closed)
+        is_true(r1._soft_closed)
+
+    def test_autoclose_rows_exhausted_plain(
+        self, connection, autoclose_row_fixture, result_fixture
+    ):
+        result = result_fixture
+
+        assert not result._soft_closed
+        assert not result.closed
+
+        read_iterator = list(result)
+        eq_(len(read_iterator), 5)
+
+        assert result._soft_closed
+        assert not result.closed
+
+        result.close()
+        assert result.closed
+
+    def test_result_ctxmanager(
+        self, connection, autoclose_row_fixture, result_fixture
+    ):
+        """test #8710"""
+
+        result = result_fixture
+
+        with expect_raises_message(Exception, "hi"):
+            with result:
+                assert not result._soft_closed
+                assert not result.closed
+
+                for i, obj in enumerate(result):
+                    if i > 2:
+                        raise Exception("hi")
+
+        assert result._soft_closed
+        assert result.closed
+
 
 class KeyTargetingTest(fixtures.TablesTest):
     run_inserts = "once"
@@ -3113,6 +3197,47 @@ class AlternateCursorResultTest(fixtures.TablesTest):
         # buffer of 98, plus buffer of 99 - 89, 10 rows
         eq_(len(result.cursor_strategy._rowbuffer), 10)
 
+        for i, row in enumerate(result):
+            if i == 206:
+                break
+
+        eq_(i, 206)
+
+    def test_iterator_remains_unbroken(self, connection):
+        """test related to #8710.
+
+        demonstrate that we can't close the cursor by catching
+        GeneratorExit inside of our iteration.  Leaving the iterable
+        block using break, then picking up again, would be directly
+        impacted by this.  So this provides a clear rationale for
+        providing context manager support for result objects.
+
+        """
+        table = self.tables.test
+
+        connection.execute(
+            table.insert(),
+            [{"x": i, "y": "t_%d" % i} for i in range(15, 250)],
+        )
+
+        result = connection.execute(table.select())
+        result = result.yield_per(100)
+        for i, row in enumerate(result):
+            if i == 188:
+                # this will raise GeneratorExit inside the iterator.
+                # so we can't close the DBAPI cursor here, we have plenty
+                # more rows to yield
+                break
+
+        eq_(i, 188)
+
+        # demonstrate getting more rows
+        for i, row in enumerate(result, 188):
+            if i == 206:
+                break
+
+        eq_(i, 206)
+
     @testing.combinations(True, False, argnames="close_on_init")
     @testing.combinations(
         "fetchone", "fetchmany", "fetchall", argnames="fetch_style"