Skip to content

A close() during a page request leaves the result set fetching #1196

Closed

Description

@laughingman7743

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

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions