fix: roll back the connection when a chunked read_sql_query fails

read_sql_query() wraps both of its return paths in a try/except that calls
con.rollback() before re-raising. For chunksize=None that works, but the
chunked path returns a generator, so nothing inside _iterate_results() runs
until the caller iterates -- long after the except clause has been unwound.
A failing statement therefore left the connection in an aborted-transaction
state, and every later query on that connection failed with a confusing
secondary error (e.g. PostgreSQL's "current transaction is aborted").

Move the rollback into _iterate_results() so both paths behave identically.
GeneratorExit is not an Exception, so abandoning the iterator early still
does not trigger a rollback.

Author: hsusul

awswrangler/_databases.py

@@ -211,30 +211,37 @@ def _iterate_results(
     timestamp_as_object: bool,
     dtype_backend: Literal["numpy_nullable", "pyarrow"],
 ) -> Iterator[pd.DataFrame]:
-    with con.cursor() as cursor:
-        cursor.execute(*cursor_args)
-        if _oracledb_found:
-            decimal_dtypes = oracle.detect_oracle_decimal_datatype(cursor)
-            _logger.debug("steporig: %s", dtype)
-            if decimal_dtypes and dtype is not None:
-                dtype = dict(list(decimal_dtypes.items()) + list(dtype.items()))
-            elif decimal_dtypes:
-                dtype = decimal_dtypes
-
-        cols_names = _get_cols_names(cursor.description)
-        while True:
-            records = cursor.fetchmany(chunksize)
-            if not records:
-                break
-            yield _records2df(
-                records=records,
-                cols_names=cols_names,
-                index=index_col,
-                safe=safe,
-                dtype=dtype,
-                timestamp_as_object=timestamp_as_object,
-                dtype_backend=dtype_backend,
-            )
+    # This generator runs lazily, so the caller's `try` block is already unwound by the time
+    # the statement is executed. The rollback must therefore happen here.
+    try:
+        with con.cursor() as cursor:
+            cursor.execute(*cursor_args)
+            if _oracledb_found:
+                decimal_dtypes = oracle.detect_oracle_decimal_datatype(cursor)
+                _logger.debug("steporig: %s", dtype)
+                if decimal_dtypes and dtype is not None:
+                    dtype = dict(list(decimal_dtypes.items()) + list(dtype.items()))
+                elif decimal_dtypes:
+                    dtype = decimal_dtypes
+
+            cols_names = _get_cols_names(cursor.description)
+            while True:
+                records = cursor.fetchmany(chunksize)
+                if not records:
+                    break
+                yield _records2df(
+                    records=records,
+                    cols_names=cols_names,
+                    index=index_col,
+                    safe=safe,
+                    dtype=dtype,
+                    timestamp_as_object=timestamp_as_object,
+                    dtype_backend=dtype_backend,
+                )
+    except Exception as ex:
+        con.rollback()
+        _logger.error(ex)
+        raise
 
 
 def _fetch_all_results(

tests/unit/test_databases.py

@@ -0,0 +1,111 @@
+"""Unit tests for the generic database helpers in ``awswrangler._databases``.
+
+These tests use an in-memory stub connection (DB-API shaped), so they require no
+database server and no AWS credentials.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Iterator
+
+import pytest
+
+from awswrangler import _databases as _db_utils
+
+
+class _StubCursor:
+    """Minimal DB-API cursor. Raises on ``execute`` or on the Nth ``fetchmany`` call."""
+
+    description = [("col0",)]
+
+    def __init__(self, fail_on_execute: bool, fail_on_fetch: bool) -> None:
+        self._fail_on_execute = fail_on_execute
+        self._fail_on_fetch = fail_on_fetch
+        self._fetch_calls = 0
+        self.closed = False
+
+    def __enter__(self) -> "_StubCursor":
+        return self
+
+    def __exit__(self, *args: object) -> None:
+        self.closed = True
+
+    def execute(self, *args: Any, **kwargs: Any) -> None:
+        if self._fail_on_execute:
+            raise RuntimeError("syntax error at or near ...")
+
+    def fetchall(self) -> list[tuple[Any, ...]]:
+        return [(1,), (2,)]
+
+    def fetchmany(self, size: int) -> list[tuple[Any, ...]]:
+        self._fetch_calls += 1
+        if self._fetch_calls == 1:
+            return [(1,)]
+        if self._fail_on_fetch:
+            raise RuntimeError("connection reset by peer")
+        return []
+
+
+class _StubConnection:
+    def __init__(self, fail_on_execute: bool = False, fail_on_fetch: bool = False) -> None:
+        self._fail_on_execute = fail_on_execute
+        self._fail_on_fetch = fail_on_fetch
+        self.rollback_count = 0
+        self.cursors: list[_StubCursor] = []
+
+    def cursor(self) -> _StubCursor:
+        cursor = _StubCursor(fail_on_execute=self._fail_on_execute, fail_on_fetch=self._fail_on_fetch)
+        self.cursors.append(cursor)
+        return cursor
+
+    def rollback(self) -> None:
+        self.rollback_count += 1
+
+
+@pytest.mark.parametrize("chunksize", [None, 1])
+def test_read_sql_query_rolls_back_on_execute_error(chunksize: int | None) -> None:
+    con = _StubConnection(fail_on_execute=True)
+
+    with pytest.raises(RuntimeError):
+        result = _db_utils.read_sql_query("SELECT 1", con=con, chunksize=chunksize)
+        if chunksize is not None:
+            list(result)
+
+    assert con.rollback_count == 1
+    assert all(cursor.closed for cursor in con.cursors)
+
+
+def test_read_sql_query_chunked_rolls_back_on_fetch_error() -> None:
+    con = _StubConnection(fail_on_fetch=True)
+    iterator: Iterator[Any] = _db_utils.read_sql_query("SELECT 1", con=con, chunksize=1)
+
+    # The first chunk is produced normally; the failure happens mid-iteration.
+    next(iterator)
+    assert con.rollback_count == 0
+
+    with pytest.raises(RuntimeError):
+        next(iterator)
+
+    assert con.rollback_count == 1
+    assert all(cursor.closed for cursor in con.cursors)
+
+
+def test_read_sql_query_chunked_does_not_roll_back_on_success() -> None:
+    con = _StubConnection()
+
+    chunks = list(_db_utils.read_sql_query("SELECT 1", con=con, chunksize=1))
+
+    assert len(chunks) == 1
+    assert chunks[0]["col0"].to_list() == [1]
+    assert con.rollback_count == 0
+
+
+def test_read_sql_query_chunked_does_not_roll_back_on_early_exit() -> None:
+    con = _StubConnection(fail_on_fetch=True)
+    iterator: Iterator[Any] = _db_utils.read_sql_query("SELECT 1", con=con, chunksize=1)
+
+    next(iterator)
+    # Abandoning the iterator throws GeneratorExit into it, which is not an error.
+    iterator.close()  # type: ignore[union-attr]
+
+    assert con.rollback_count == 0