Share the buffer refill between fetchone() and the bulk fetches of the result sets

[laughingman7743]

WHAT

Option 1 of #1193: one refill condition shared by fetchone() and the bulk fetches.

• AthenaResultSet._fetch() and AthenaAioResultSet._async_fetch() return without rows when no NextToken remains, instead of raising ProgrammingError("NextToken is none or empty."). In every result set, _fetch() now means "read the next rows, if any", as it already did for Arrow and S3FS.
• The base and asyncio fetchone() call it whenever the buffer is empty. The Arrow and S3FS fetchone() overrides only differed in that condition, so they are removed.
• _take_rows() and _async_take_rows() call _fetch() / _async_fetch() directly when the buffer is empty and take every row in bulk. Before, they took the first row of each page or batch through fetchone().

Public behavior is unchanged: the rows, rownumber, and the GetQueryResults requests stay the same. A page that comes back empty still ends the fetch after one request, as fetchone() did before.
The private contract changes in one way: a subclass whose fetchone() doesn't use the buffer now has to override _take_rows() as well. The pandas and Polars result sets already do; no other PyAthena result set has such a fetchone(). So the test of that former contract (_FetchoneResultSet) is removed.
Released 4.1.0 still looped fetchone() in fetchmany() and fetchall(), so this is a release-note candidate for 4.2.0: a custom result set subclass that overrides only fetchone() must also override _take_rows(). dbt-athena, which builds AthenaResultSet directly, overrides neither method and doesn't call _fetch().

WHY

Closes #1193. Since #1187, the same fetchone() body existed in the base, Arrow, S3FS, and asyncio result sets, and the bulk take used fetchone() only to trigger a refill. The maintainer chose option 1 of #1193 over a new _fetch_next() hook.

TEST

Tested at 0100221 unless noted; the head 6cd9b59 changes only two docstrings after it (lint and the base/aio result set unit tests passed there).

• just format and just lint: passed (ruff, mypy, cfn-lint, license headers).
• Unit tests: uv run --env-file .env pytest tests/pyathena/test_result_set.py tests/pyathena/aio/test_result_set.py tests/pyathena/s3fs/test_result_set.py tests/pyathena/arrow/test_result_set.py tests/pyathena/pandas/test_result_set.py tests/pyathena/polars/test_result_set.py: 354 passed, 10 skipped.
  • New: base/aio fetchone() across pages with the request count (the last call reaches _fetch() with no page left, which raised ProgrammingError before), aio fetches after close() and on a failed page request, and Arrow/S3FS refills across batches up to the end of the results.
• Against Athena: uv run --env-file .env pytest -n 1 -k "fetch or iter" on the default, Arrow, and S3FS cursor tests (sync, async, and aio): 74 passed at 63646b5. Later commits change only docstrings and unit tests.
• A local micro-benchmark on an Apple M-series laptop with Python 3.15 (AthenaResultSet over 300 stubbed GetQueryResults pages × 1,000 rows, best of 5) showed no measurable change: fetchone() loop 0.199 → 0.195 s, fetchmany(1) loop 0.316 → 0.294 s, fetchall() 0.177 → 0.184 s.
• Not run locally: the full just test pyathena, just test sqla, and just test sqla-async. The AWS CI runs them when the PR is marked Ready.

🤖 Generated with Claude Code

[laughingman7743]

Self-review round one (implementation behavior): FINDINGS, repaired

Scope: git diff 06224494788043a4887afcb112a96c94bbcb02f1..63646b5b78e2f7cea6565c7b5b49a923f614fd3c, full pass over all 8 files. It covered the base/aio _fetch(), fetchone(), and _take_rows(); the removed Arrow/S3FS fetchone(); every AthenaResultSet subclass (dict, aio dict through the MRO, pandas/Polars overrides); and the cursors that call these methods. It also ran /code-review (high) and /simplify (reuse, simplification, efficiency, altitude).

Verified unchanged: the rows, rownumber, and the GetQueryResults request count. A page that comes back empty still ends the fetch after one request, as fetchone() did before. After close(), _next_token is None and the buffer is empty, so no request is made. Arrow and S3FS _fetch() already returned without rows at the end, so they inherit the base fetchone() unchanged. A micro-benchmark showed no measurable change.

Repairs in 0100221:

1. _fetch() / _async_fetch() docstrings: "fetches nothing when … the result set is closed" holds only when close() comes before the call, and "another thread" is the wrong model for asyncio. Both now say "when no page remains, as after close()" and "or the result set is closed", the wording _get_query_results() uses.
2. Removed test_fetch_without_next_page (sync/aio). test_fetchone_across_pages covers the same path through the public API: its fourth fetchone() reaches _fetch() with no token, and call_count == 2 is asserted. Also dropped the duplicate stream.closed assertion in the S3FS refill test, which test_reader_closes_at_end_of_results already checks.

Deferred with reasons:

• A subclass that overrides only fetchone() now needs _take_rows() too. This is the agreed private contract of Share the buffer refill between fetchone() and the bulk fetches of the result sets #1193 and is documented on _take_rows(); see round two for the release-note question.
• Rows taken before a failed refill are counted and lost with the exception. This predates the PR: released 4.1.0 looped fetchone() with the same result, and the sync test is from Take buffered rows in bulk in fetchmany() and fetchall() #1187.
• An empty refill (an empty GetQueryResults page with a NextToken, or an empty Arrow batch) ends the fetch. This is also pre-existing: the old fetchone() returned None there too.
• The pandas/Polars fetchone()/_take_rows() duplicates predate the PR. Sharing them needs a class restructuring, which is out of scope.

[laughingman7743]

Self-review round two (claims, compatibility, operations): FINDINGS, repaired in the PR description

Scope: git diff 06224494788043a4887afcb112a96c94bbcb02f1..0100221287e8329c273967a3b1923c2ce9dce287, full claim audit of the PR body, the commit messages, and the changed docstrings.

Claims checked:

• "Rows, rownumber, and GetQueryResults requests stay the same": checked by the stubbed request counts in the base/aio across-pages tests, by the Arrow/S3FS refill tests, and by 74 live fetch tests at 63646b5. The source changes since then are docstrings only.
• "No other PyAthena result set has a fetchone() that ignores the buffer": a grep of def fetchone under pyathena/ finds only base, aio, pandas, and Polars on result sets. Pandas and Polars both override _take_rows().
• _fetch() docstrings: no page remains after close() because close() clears _next_token (pyathena/result_set.py:880). The "closed" ProgrammingError comes from _get_query_results() / _async_get_query_results().
• No docs, comments, or tests still mention the removed "NextToken is none or empty." error or the old route through fetchone().

Existing callers: dbt-athena (dbt-labs/dbt-adapters, connections_legacy.py:157) builds AthenaResultSet directly and doesn't override fetchone() or call _fetch(). Its connections.py uses its own result set class.

AWS operator: no change to retries or request counts. An exhausted or closed result set now calls _fetch(), which returns before any request.

Corrections to the PR description:

1. Added a 4.2.0 release-note candidate. Released 4.1.0 still looped fetchone() in fetchmany()/fetchall(), so a custom subclass that overrides only fetchone() must now also override _take_rows().
2. "approved option 1 in Share the buffer refill between fetchone() and the bulk fetches of the result sets #1193" now reads "chose option 1". The decision was made with the maintainer, not recorded on the issue.
3. Updated the TEST section: the tested commit, 354 unit tests after the round-one test removals, the commit the live tests ran at, and the benchmark environment.

Limits: the full just test pyathena / sqla / sqla-async suites run in the AWS CI once the PR is Ready. The micro-benchmark used stubbed pages on a laptop, not AWS traffic.

[laughingman7743]

Independent review (relayed), static only: FINDINGS

Reviewer: Codex CLI 0.162.1 (codex exec -s read-only), model gpt-6-astra, session 01a14031-fbf3-7883-941a-4fbeed65dd05. Scope: git diff 06224494788043a4887afcb112a96c94bbcb02f1..0100221287e8329c273967a3b1923c2ce9dce287 on a detached snapshot at 0100221. The prompt left out the PR description and the prior findings, and allowed no edits, builds, tests, GitHub, or network. The snapshot was unchanged afterwards.

Coverage: base, dict, and asyncio result sets (including the AthenaAioDictResultSet MRO) and all three fetch methods; the Arrow/S3FS refill; the pandas/Polars bulk overrides; the sync, futures-based, and aio cursor callers; request counts, failure accounting, rownumber, close and concurrent close, docstrings, and tests.

Findings:

1. P2, pre-existing: pyathena/aio/result_set.py:158 and pyathena/result_set.py:475. If close() completes during a page request and the response carries another NextToken, the response restores that token. The cleared metadata keeps rows out of the buffer, but the next fetch requests again and raises ProgrammingError for the missing query ID. The reviewer confirmed that the base commit has the same defect.
   → Deferred as pre-existing and outside Share the buffer refill between fetchone() and the bulk fetches of the result sets #1193. The old fetchone() hit the same request through its _next_token check. A follow-up issue is proposed to the maintainer.
2. P3, introduced: the Raises: ProgrammingError text on _fetch() / _async_fetch() said a closed result set raises, but close(); _fetch() returns without error because the token is cleared.
   → Verified and repaired in 6cd9b59 (reply below).

The reviewer found no other functional regression and judged that the tests exercise page/batch refills, exhaustion, request counts, and partial-consumption accounting.

[laughingman7743]

Repair of finding 2: 6cd9b59 (docstrings only, base and aio).

Raises: ProgrammingError now reads "If the next page is requested and the query ID is missing, the query has not succeeded, or the result set is closed."

• Round-one perspective: no code changed. just lint passed, and tests/pyathena/test_result_set.py plus tests/pyathena/aio/test_result_set.py passed (27).
• Round-two perspective: the claim matches the code. The error comes only from _build_get_query_results_request() or the closed check in _get_query_results() / _async_get_query_results(), and both run only when a _next_token remains.
• Independent follow-up: Codex (gpt-6-astra, codex exec -s read-only, session 01a14038-971d-7491-896b-0f610f70ec86) reviewed git diff 0100221287e8329c273967a3b1923c2ce9dce287..6cd9b5962dcc93ac668eff417e137641aef5e0c4 on a detached snapshot and returned CLEAN: the repair resolves the finding, describes both methods accurately, and adds nothing new. This was a static review.