fix: roll back the connection when a chunked read_sql_query fails

[hsusul]

Feature or Bugfix

• Bugfix

Detail

• read_sql_query(..., chunksize=N) never rolled back the connection when the query failed, because _iterate_results is a generator and its body runs after the enclosing try/except has already been unwound. The non-chunked path (chunksize=None) has always rolled back. This fixes the chunked path so both behave identically.

Relates

• No existing issue — found while auditing serial vs. chunked parity in the database read paths. Reproducible locally with a stub DB-API connection (details below); no AWS credentials or database server required.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.

---

Affected API

awswrangler._databases.read_sql_query, and therefore every public wrapper that routes through it with chunksize:

• wr.postgresql.read_sql_query / read_sql_table
• wr.mysql.read_sql_query / read_sql_table
• wr.redshift.read_sql_query / read_sql_table
• wr.sqlserver.read_sql_query / read_sql_table
• wr.oracle.read_sql_query / read_sql_table

Current behavior

read_sql_query wraps both of its return paths in a try/except that calls con.rollback() before re-raising:

    try:
        if chunksize is None:
            return _fetch_all_results(...)

        return _iterate_results(...)
    except Exception as ex:
        con.rollback()
        _logger.error(ex)
        raise

_iterate_results is a generator function, so calling it only builds a generator object — cursor.execute(...) does not run until the caller iterates, by which point the except clause has already been unwound. The rollback is therefore dead code for every chunked read: a failing statement leaves the connection in an aborted-transaction state, and the next query on that same connection fails with a confusing secondary error (on PostgreSQL/Redshift: current transaction is aborted, commands ignored until end of transaction block).

The same applies to a failure that happens part-way through iteration (e.g. a dropped connection during fetchmany).

Reproduction

No database server and no AWS credentials required — a DB-API-shaped stub is enough:

from awswrangler import _databases as db


class Cursor:
    description = [("a",)]

    def __enter__(self): return self
    def __exit__(self, *args): return None
    def execute(self, *args): raise RuntimeError("syntax error at or near ...")
    def fetchall(self): return []
    def fetchmany(self, n): return []


class Con:
    def __init__(self): self.rollbacks = 0
    def cursor(self): return Cursor()
    def rollback(self): self.rollbacks += 1


con = Con()
try:
    db.read_sql_query("SELECT 1", con=con)
except RuntimeError:
    pass
print("non-chunked rollbacks:", con.rollbacks)   # 1

con = Con()
try:
    list(db.read_sql_query("SELECT 1", con=con, chunksize=10))
except RuntimeError:
    pass
print("chunked rollbacks:", con.rollbacks)       # 0  <-- expected 1

Output on main (609db3c):

non-chunked rollbacks: 1
chunked     rollbacks: 0 (expected 1)

Corrected behavior

Both paths roll back the connection before re-raising, so a failed chunked read leaves the connection in the same usable state as a failed non-chunked read.

Root cause

Lazy generator evaluation: the try block in read_sql_query only covers the construction of the generator, not its execution.

Implementation

Move the rollback into _iterate_results, wrapping its body — this covers both the initial cursor.execute and any failure during fetchmany/frame construction. read_sql_query is otherwise unchanged, and the non-chunked path is untouched.

GeneratorExit is a BaseException, not an Exception, so abandoning the iterator early (break, .close(), garbage collection) still does not trigger a rollback — normal early exit is not an error.

Regression tests

New tests/unit/test_databases.py (pure unit tests, stub connection, no server and no AWS credentials):

• test_read_sql_query_rolls_back_on_execute_error[None] / [1] — parametrized over chunksize, asserts identical rollback behavior on both paths and that the cursor context manager is exited. The [1] case fails on main.
• test_read_sql_query_chunked_rolls_back_on_fetch_error — failure part-way through iteration (first chunk yielded successfully, second fetchmany raises); fails on main.
• test_read_sql_query_chunked_does_not_roll_back_on_success — no rollback on the happy path, and the yielded frame carries the expected values.
• test_read_sql_query_chunked_does_not_roll_back_on_early_exit — .close() on the iterator does not roll back.

Before the fix: 2 failed, 3 passed. After: 5 passed.

Validation

Run locally on macOS / Python 3.13.5 (pandas 3.0.1, pyarrow 25.0.0, moto 5.2.2):

Check	Command	Result
New regression tests	pytest tests/unit/test_databases.py	5 passed
Credential-free unit suites	pytest tests/unit/{test_databases,test_moto,test_utils,test_data_types,test_sql_params_formatter,test_s3_vectors_mocked,test_metadata}.py	142 passed, 1 pre-existing failure
Format	ruff format --check .	281 files already formatted
Lint	ruff check .	All checks passed
Types	mypy awswrangler	19 errors — identical count and files on unmodified main; none in _databases.py
Docs	doc8 --ignore-path docs/source/stubs --max-line-length 120 docs/source	exit 0
Lock	uv lock --check	up to date
Build	uv build	sdist + wheel built
Whitespace	git diff --check	clean

The one failure, test_s3_vectors_mocked.py::test_delete_vector_index_arn_with_name_raises, reproduces identically on unmodified main and is a local-environment artifact (botocore.exceptions.NoRegionError: You must specify a region) — unrelated to this change.

Not run: every suite that needs live AWS infrastructure or a real database server — tests/unit/test_postgresql.py, test_mysql.py, test_redshift.py, test_sqlserver.py, test_oracle.py (these are the integration tests that exercise the touched code path against real engines), plus the Athena/Glue/S3/DynamoDB/OpenSearch/Timestream/QuickSight/EMR suites and tests/load. tests/unit/test_neptune_parsing.py was skipped because gremlin_python is not installed locally. The new tests are the credential-free substitute for the affected path.

Compatibility and ownership

• No public API, signature, or return-type change.
• No dependency or lockfile change.
• Purely additive error handling: the success path is byte-for-byte identical, dtypes and frame contents are unaffected.
• The connection stays caller-owned — it is neither closed nor reconfigured; only rollback() is called, mirroring what the non-chunked path has always done.
• Callers that already wrap chunked reads in their own try/rollback see a harmless second rollback on an already-clean transaction.