Merge "ensure soft_close occurs for fetchmany with server side cursor" into main

3 files changed, 154 insertions, 2 deletions

diff --git a/test/base/test_result.py b/test/base/test_result.py
index d94602203..8c9eb398e 100644
--- a/test/base/test_result.py
+++ b/test/base/test_result.py
@@ -484,7 +484,12 @@ class ResultTest(fixtures.TestBase):
         row = result.first()
         eq_(row, (1, 1, 1))
 
-        eq_(result.all(), [])
+        # note this is a behavior change in 1.4.27 due to
+        # adding a real result.close() to Result, previously this would
+        # return an empty list.  this is already the
+        # behavior with CursorResult, but was mis-implemented for
+        # other non-cursor result sets.
+        assert_raises(exc.ResourceClosedError, result.all)
 
     def test_one_unique(self):
         # assert that one() counts rows after uniqueness has been applied.
@@ -597,7 +602,12 @@ class ResultTest(fixtures.TestBase):
 
         eq_(result.scalar(), 1)
 
-        eq_(result.all(), [])
+        # note this is a behavior change in 1.4.27 due to
+        # adding a real result.close() to Result, previously this would
+        # return an empty list.  this is already the
+        # behavior with CursorResult, but was mis-implemented for
+        # other non-cursor result sets.
+        assert_raises(exc.ResourceClosedError, result.all)
 
     def test_partition(self):
         result = self._fixture()
diff --git a/test/orm/test_query.py b/test/orm/test_query.py
index b8269d76d..98be71f34 100644
--- a/test/orm/test_query.py
+++ b/test/orm/test_query.py
@@ -73,6 +73,7 @@ from sqlalchemy.testing import mock
 from sqlalchemy.testing.assertions import assert_raises
 from sqlalchemy.testing.assertions import assert_raises_message
 from sqlalchemy.testing.assertions import eq_
+from sqlalchemy.testing.assertions import expect_raises
 from sqlalchemy.testing.assertions import expect_warnings
 from sqlalchemy.testing.assertions import is_not_none
 from sqlalchemy.testing.assertsql import CompiledSQL
@@ -5270,6 +5271,8 @@ class YieldTest(_fixtures.FixtureTest):
     run_setup_mappers = "each"
     run_inserts = "each"
 
+    __backend__ = True
+
     def _eagerload_mappings(self, addresses_lazy=True, user_lazy=True):
         User, Address = self.classes("User", "Address")
         users, addresses = self.tables("users", "addresses")
@@ -5312,6 +5315,83 @@ class YieldTest(_fixtures.FixtureTest):
         except StopIteration:
             pass
 
+    def test_we_can_close_cursor(self):
+        """test new usecase close() added along with #7274"""
+        self._eagerload_mappings()
+
+        User = self.classes.User
+
+        sess = fixture_session()
+
+        stmt = select(User).execution_options(yield_per=15)
+        result = sess.execute(stmt)
+
+        with mock.patch.object(result.raw, "_soft_close") as mock_close:
+            two_results = result.fetchmany(2)
+            eq_(len(two_results), 2)
+
+            eq_(mock_close.mock_calls, [])
+
+            result.close()
+
+            eq_(mock_close.mock_calls, [mock.call(hard=True)])
+
+            with expect_raises(sa.exc.ResourceClosedError):
+                result.fetchmany(10)
+
+            with expect_raises(sa.exc.ResourceClosedError):
+                result.fetchone()
+
+            with expect_raises(sa.exc.ResourceClosedError):
+                result.all()
+
+        result.close()
+
+    @testing.combinations("fetchmany", "fetchone", "fetchall")
+    def test_cursor_is_closed_on_exhausted(self, fetch_method):
+        """test #7274"""
+        self._eagerload_mappings()
+
+        User = self.classes.User
+
+        sess = fixture_session()
+
+        stmt = select(User).execution_options(yield_per=15)
+        result = sess.execute(stmt)
+
+        with mock.patch.object(result.raw, "_soft_close") as mock_close:
+            # call assertions are implementation specific.
+            # test needs that _soft_close called at least once and without
+            # the hard=True flag
+            if fetch_method == "fetchmany":
+                while True:
+                    buf = result.fetchmany(2)
+                    if not buf:
+                        break
+                eq_(mock_close.mock_calls, [mock.call()])
+            elif fetch_method == "fetchall":
+                eq_(len(result.all()), 4)
+                eq_(
+                    mock_close.mock_calls, [mock.call(), mock.call(hard=False)]
+                )
+            elif fetch_method == "fetchone":
+                while True:
+                    row = result.fetchone()
+                    if row is None:
+                        break
+                eq_(
+                    mock_close.mock_calls, [mock.call(), mock.call(hard=False)]
+                )
+            else:
+                assert False
+
+        # soft closed, we can still get an empty result
+        eq_(result.all(), [])
+
+        # real closed
+        result.close()
+        assert_raises(sa.exc.ResourceClosedError, result.all)
+
     def test_yield_per_and_execution_options_legacy(self):
         self._eagerload_mappings()
 
diff --git a/test/sql/test_resultset.py b/test/sql/test_resultset.py
index 86aa15f26..346cb3d58 100644
--- a/test/sql/test_resultset.py
+++ b/test/sql/test_resultset.py
@@ -2736,6 +2736,68 @@ class AlternateCursorResultTest(fixtures.TablesTest):
         # buffer of 98, plus buffer of 99 - 89, 10 rows
         eq_(len(result.cursor_strategy._rowbuffer), 10)
 
+    @testing.combinations(True, False, argnames="close_on_init")
+    @testing.combinations(
+        "fetchone", "fetchmany", "fetchall", argnames="fetch_style"
+    )
+    def test_buffered_fetch_auto_soft_close(
+        self, connection, close_on_init, fetch_style
+    ):
+        """test #7274"""
+
+        table = self.tables.test
+
+        connection.execute(
+            table.insert(),
+            [{"x": i, "y": "t_%d" % i} for i in range(15, 30)],
+        )
+
+        result = connection.execute(table.select().limit(15))
+        assert isinstance(result.cursor_strategy, _cursor.CursorFetchStrategy)
+
+        if close_on_init:
+            # close_on_init - the initial buffering will exhaust the cursor,
+            # should soft close immediately
+            result = result.yield_per(30)
+        else:
+            # not close_on_init - soft close will occur after fetching an
+            # empty buffer
+            result = result.yield_per(5)
+        assert isinstance(
+            result.cursor_strategy, _cursor.BufferedRowCursorFetchStrategy
+        )
+
+        with mock.patch.object(result, "_soft_close") as soft_close:
+            if fetch_style == "fetchone":
+                while True:
+                    row = result.fetchone()
+
+                    if row:
+                        eq_(soft_close.mock_calls, [])
+                    else:
+                        # fetchone() is also used by first(), scalar()
+                        # and one() which want to embed a hard close in one
+                        # step
+                        eq_(soft_close.mock_calls, [mock.call(hard=False)])
+                        break
+            elif fetch_style == "fetchmany":
+                while True:
+                    rows = result.fetchmany(5)
+
+                    if rows:
+                        eq_(soft_close.mock_calls, [])
+                    else:
+                        eq_(soft_close.mock_calls, [mock.call()])
+                        break
+            elif fetch_style == "fetchall":
+                rows = result.fetchall()
+
+                eq_(soft_close.mock_calls, [mock.call()])
+            else:
+                assert False
+
+        result.close()
+
     def test_buffered_fetchmany_yield_per_all(self, connection):
         table = self.tables.test