fix(spanner): buffer returned rows in autocommit DML statements (#18152)

### Problem
In autocommit mode (`connection.autocommit = True`), `cursor.execute()`
runs within a `run_in_transaction` callback that commits and terminates
the underlying Spanner transaction immediately upon return.
Because `PeekIterator` only prefetched the first row from the streaming
`StreamedResultSet`, subsequent rows from multi-row DML statements with
returning clauses (e.g. `INSERT ... THEN RETURN` or `UPDATE ... THEN
RETURN`) were not fetched before transaction closure. Attempting to
consume subsequent rows post-commit caused stream errors, resulting in
silently dropped rows for multi-row returning queries.
### Solution
- Introduced `BufferedIterator` in `google.cloud.spanner_dbapi.utils` to
eagerly drain and buffer returned rows from `StreamedResultSet` into
memory upon instantiation while converting row lists to tuples per
DB-API v2 (PEP 249) expectations.
- Updated `_do_execute_update_in_autocommit` in
`google.cloud.spanner_dbapi.cursor` to use
`BufferedIterator(self._result_set)`, ensuring all rows and result set
stats are captured before the transaction commits.
- Added comprehensive unit tests in `test_utils.py` and `test_cursor.py`
verifying eager draining, list-to-tuple conversions, and multi-row
post-commit fetch behavior.

Thank you for opening a Pull Request! Before submitting your PR, there
are a few things you can do to make sure it goes smoothly:
- [ ] Make sure to open an issue as a
[bug/issue](https://github.com/googleapis/google-cloud-python/issues)
before writing your code! That way we can discuss the change, evaluate
designs, and agree on the general idea
- [ ] Ensure the tests and linter pass
- [ ] Code coverage does not decrease (if any source code was changed)
- [ ] Appropriate docs were updated (if necessary)

Fixes #<issue_number_goes_here> 🦕

Author: sakthivelmanii

packages/google-cloud-spanner/google/cloud/spanner_dbapi/cursor.py

@@ -49,7 +49,11 @@
     StatementType,
 )
 from google.cloud.spanner_dbapi.transaction_helper import CursorStatementType
-from google.cloud.spanner_dbapi.utils import PeekIterator, StreamedManyResultSets
+from google.cloud.spanner_dbapi.utils import (
+    BufferedIterator,
+    PeekIterator,
+    StreamedManyResultSets,
+)
 from google.cloud.spanner_v1 import RequestOptions
 from google.cloud.spanner_v1.merged_result_set import MergedResultSet
 
@@ -234,7 +238,7 @@ def _do_execute_update_in_autocommit(self, transaction, sql, params):
             param_types=get_param_types(params),
             last_statement=True,
         )
-        self._itr = PeekIterator(self._result_set)
+        self._itr = BufferedIterator(self._result_set)
         self._row_count = None
 
     def _batch_DDLs(self, sql):

packages/google-cloud-spanner/google/cloud/spanner_dbapi/utils.py

@@ -17,6 +17,36 @@
 re_UNICODE_POINTS = re.compile(r"([^\s]*[\u0080-\uFFFF]+[^\s]*)")
 
 
+class BufferedIterator:
+    """
+    An eager, in-memory iterator that consumes and buffers all rows from a
+    source stream upon instantiation.
+
+    Used when the underlying database transaction is committed and closed before
+    the caller fetches results (e.g. autocommit DML with THEN RETURN / RETURNING).
+    If a row is an instance of list, it is converted to a tuple to conform
+    with DBAPI v2's sequence expectations.
+
+    :type source: iterable
+    :param source: A source iterable/stream of rows.
+    """
+
+    def __init__(self, source):
+        self._rows = [
+            tuple(row) if isinstance(row, list) else row for row in (source or ())
+        ]
+        self._itr = iter(self._rows)
+
+    def __next__(self):
+        return next(self._itr)
+
+    def __iter__(self):
+        return self
+
+    def __len__(self):
+        return len(self._rows)
+
+
 class PeekIterator:
     """
     Peek at the first element out of an iterator for the sake of operations

packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_cursor.py

@@ -123,6 +123,58 @@ def test_do_execute_update(self):
         self.assertEqual(cursor._result_set, result_set)
         self.assertEqual(cursor.rowcount, 1234)
 
+    def test_do_execute_update_in_autocommit_eagerly_buffers_returning_rows(self):
+        from google.cloud.spanner_v1 import ResultSetStats, Type, TypeCode
+        from google.cloud.spanner_v1.types import ResultSetMetadata, StructType
+
+        connection = self._make_connection(self.INSTANCE, mock.MagicMock())
+        connection.autocommit = True
+        cursor = self._make_one(connection)
+
+        tx_active = True
+
+        def streaming_generator():
+            if not tx_active:
+                raise ValueError(
+                    "Transaction has already been committed or rolled back"
+                )
+            yield [1, "Alice"]
+            if not tx_active:
+                raise ValueError(
+                    "Transaction has already been committed or rolled back"
+                )
+            yield [2, "Bob"]
+
+        field1 = StructType.Field(name="id", type=Type(code=TypeCode.INT64))
+        field2 = StructType.Field(name="name", type=Type(code=TypeCode.STRING))
+        row_type = StructType(fields=[field1, field2])
+        metadata = ResultSetMetadata(row_type=row_type)
+
+        mock_result_set = mock.MagicMock()
+        mock_result_set.__iter__.side_effect = streaming_generator
+        mock_result_set.metadata = metadata
+        mock_result_set.stats = ResultSetStats(row_count_exact=2)
+
+        def fake_run_in_transaction(func, *args, **kwargs):
+            nonlocal tx_active
+            mock_tx = mock.MagicMock()
+            mock_tx.execute_sql.return_value = mock_result_set
+            result = func(mock_tx, *args, **kwargs)
+            tx_active = False
+            return result
+
+        connection.database.run_in_transaction = fake_run_in_transaction
+
+        sql = "INSERT INTO table (id, name) VALUES (1, 'Alice'), (2, 'Bob') THEN RETURN id, name"
+        cursor.execute(sql)
+
+        self.assertIsNotNone(cursor.description)
+        self.assertEqual(len(cursor.description), 2)
+        self.assertEqual(cursor.description[0].name, "id")
+        self.assertEqual(cursor.description[1].name, "name")
+        self.assertEqual(cursor.rowcount, 2)
+        self.assertEqual(cursor.fetchall(), [(1, "Alice"), (2, "Bob")])
+
     def test_do_batch_update(self):
         from google.cloud.spanner_dbapi import connect
         from google.cloud.spanner_v1.param_types import INT64

packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_utils.py

@@ -12,15 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import sys
 import unittest
 
 
 class TestUtils(unittest.TestCase):
-    skip_condition = sys.version_info[0] < 3
-    skip_message = "Subtests are not supported in Python 2"
-
-    @unittest.skipIf(skip_condition, skip_message)
     def test_PeekIterator(self):
         from google.cloud.spanner_dbapi.utils import PeekIterator
 
@@ -38,7 +33,6 @@ def test_PeekIterator(self):
                 actual = list(pitr)
                 self.assertEqual(actual, expected)
 
-    @unittest.skipIf(skip_condition, "Python 2 has an outdated iterator definition")
     def test_peekIterator_list_rows_converted_to_tuples(self):
         from google.cloud.spanner_dbapi.utils import PeekIterator
 
@@ -60,7 +54,6 @@ def test_peekIterator_list_rows_converted_to_tuples(self):
         pit = PeekIterator([("Clark", "Kent")])
         self.assertEqual(next(pit), ("Clark", "Kent"))
 
-    @unittest.skipIf(skip_condition, "Python 2 has an outdated iterator definition")
     def test_peekIterator_nonlist_rows_unconverted(self):
         from google.cloud.spanner_dbapi.utils import PeekIterator
 
@@ -69,7 +62,6 @@ def test_peekIterator_nonlist_rows_unconverted(self):
         want = ["a", "b", "c", "d", "e"]
         self.assertEqual(got, want, "Values should be returned unchanged")
 
-    @unittest.skipIf(skip_condition, skip_message)
     def test_backtick_unicode(self):
         from google.cloud.spanner_dbapi.utils import backtick_unicode
 
@@ -85,7 +77,6 @@ def test_backtick_unicode(self):
                 got = backtick_unicode(sql)
                 self.assertEqual(got, want)
 
-    @unittest.skipIf(skip_condition, skip_message)
     def test_StreamedManyResultSets(self):
         from google.cloud.spanner_dbapi.utils import StreamedManyResultSets
 
@@ -100,3 +91,45 @@ def test_StreamedManyResultSets(self):
                 stream_result._iterators.append(data_in)
                 actual = list(stream_result)
                 self.assertEqual(actual, expected)
+
+    def test_BufferedIterator(self):
+        from google.cloud.spanner_dbapi.utils import BufferedIterator
+
+        cases = [
+            ("list_of_lists", [["a", 1], ["b", 2]], [("a", 1), ("b", 2)]),
+            ("iter_of_lists", iter([["a", 1], ["b", 2]]), [("a", 1), ("b", 2)]),
+            ("list_of_tuples", [("a", 1), ("b", 2)], [("a", 1), ("b", 2)]),
+            ("iter_of_tuples", iter([("a", 1), ("b", 2)]), [("a", 1), ("b", 2)]),
+            ("empty_list", [], []),
+            ("empty_tuple", (), []),
+            ("none_source", None, []),
+        ]
+
+        for name, data_in, expected in cases:
+            with self.subTest(name=name):
+                bitr = BufferedIterator(data_in)
+                self.assertEqual(len(bitr), len(expected))
+                actual = list(bitr)
+                self.assertEqual(actual, expected)
+                with self.assertRaises(StopIteration):
+                    next(bitr)
+
+    def test_BufferedIterator_eagerly_drains_stream(self):
+        from google.cloud.spanner_dbapi.utils import BufferedIterator
+
+        drained = False
+
+        def generator():
+            nonlocal drained
+            yield ["first", 1]
+            yield ["second", 2]
+            drained = True
+
+        bitr = BufferedIterator(generator())
+        # The generator must be completely drained during __init__
+        self.assertTrue(drained)
+        self.assertEqual(len(bitr), 2)
+        self.assertEqual(next(bitr), ("first", 1))
+        self.assertEqual(next(bitr), ("second", 2))
+        with self.assertRaises(StopIteration):
+            next(bitr)