Repository navigation
A close() during a page request leaves the result set fetching #1196
Description
Problem
AthenaResultSet.close() clears _next_token and the buffer, so a fetch after close() returns None (or []) without a request.
If close() runs while _fetch() is waiting for a GetQueryResults response, _fetch() then stores the response's NextToken in _next_token (pyathena/result_set.py, rows, self._next_token = self._parse_result_rows(response)).
The cleared metadata keeps the response's rows out of the buffer, but the restored token makes the next fetch request another page, which raises ProgrammingError("QueryExecutionId is none or empty.") instead of returning None.
AthenaAioResultSet._async_fetch() has the same assignment after await self._async_get_query_results(...), where another task can call close() during the await.
Expected: a closed result set fetches nothing, whether close() came before or during a page request.
Reproduction
No AWS calls; GetQueryResults is stubbed and closes the result set during the second request, as another thread would:
from unittest.mock import MagicMock
from pyathena.converter import DefaultTypeConverter
from pyathena.model import AthenaQueryExecution
from pyathena.result_set import AthenaResultSet
from pyathena.util import RetryConfig
def page(values, next_token=None, metadata=False):
response = {"ResultSet": {"Rows": [{"Data": [{"VarCharValue": v}]} for v in values]}}
if metadata:
response["ResultSet"]["ResultSetMetadata"] = {"ColumnInfo": [{"Name": "a", "Type": "integer"}]}
if next_token:
response["NextToken"] = next_token
return response
responses = iter([page(["1"], "token1", metadata=True), page(["2"], "token2"), page(["3"])])
def get_query_results(**kwargs):
response = next(responses)
if kwargs.get("NextToken") == "token1":
result_set.close() # another thread closes the result set during the request
return response
connection = MagicMock()
connection.client.get_query_results.side_effect = get_query_results
result_set = AthenaResultSet(
connection=connection,
converter=DefaultTypeConverter(),
query_execution=MagicMock(state=AthenaQueryExecution.STATE_SUCCEEDED, query_id="q"),
arraysize=1,
retry_config=RetryConfig(),
)
print(result_set.fetchone()) # (1,)
print(result_set.fetchone(), result_set.is_closed, result_set._next_token) # None True token2
print(result_set.fetchone()) # ProgrammingError: QueryExecutionId is none or empty.Environment
- PyAthena master at 7dd8ee4, and the branch of Share the buffer refill between fetchone() and the bulk fetches of the result sets #1195 (the code path is the same in v4.1.0)
- Python 3.15
- Default cursor (
AthenaResultSet),AthenaDictResultSet, and the asyncioAthenaAioResultSet/AthenaAioDictResultSet
Found during the independent review of #1195; it predates that change.
Proposed fix (optional)
In _fetch() and _async_fetch(), discard the response when the result set has been closed during the request, before assigning _next_token.
Validate with unit tests in tests/pyathena/test_result_set.py and tests/pyathena/aio/test_result_set.py that close the result set from the stubbed request, as above, and check that later fetches return None / [] without another request.
Activity