diff --git a/pyathena/aio/result_set.py b/pyathena/aio/result_set.py index b55627ab2..5e27e4b44 100644 --- a/pyathena/aio/result_set.py +++ b/pyathena/aio/result_set.py @@ -143,14 +143,17 @@ async def _async_get_query_results( return cast(dict[str, Any], response) async def _async_fetch(self) -> None: - """Fetch the next page of rows into the result set. + """Fetch the next page of rows, if any, into the buffer. + + Nothing is fetched when no page remains, as after ``close()``. Raises: - ProgrammingError: If there is no next page. + ProgrammingError: If the next page is requested and the query ID is + missing, the query has not succeeded, or the result set is closed. OperationalError: If the request fails. """ if not self._next_token: - raise ProgrammingError("NextToken is none or empty.") + return response = await self._async_get_query_results(self._arraysize, self._next_token) rows, self._next_token = self._parse_result_rows(response) self._process_rows(rows) @@ -183,10 +186,10 @@ async def fetchone( # type: ignore[override] The next row (a tuple, or a dict for ``AthenaAioDictResultSet``), or None if no more rows. """ - if not self._rows and self._next_token: - await self._async_fetch() if not self._rows: - return None + await self._async_fetch() + if not self._rows: + return None row = self._rows.popleft() # As _count_rows() does, inline on the per-row path. rownumber = self._rownumber @@ -230,8 +233,8 @@ async def _async_take_rows( ) -> list[tuple[Any | None, ...] | dict[Any, Any | None]]: """Take the next rows and count them in ``rownumber``. - The buffered rows are taken in bulk. When the buffer is empty, ``fetchone()`` - takes the next row, so that it fetches the next page into the buffer. + The buffered rows are taken in bulk, and ``_async_fetch()`` refills the buffer + when it is empty. Args: size: The maximum number of rows, or None for all remaining rows. @@ -241,13 +244,11 @@ async def _async_take_rows( """ rows: list[tuple[Any | None, ...] | dict[Any, Any | None]] = [] while size is None or len(rows) < size: - if self._rows: - self._take_buffered_rows(rows, size) - else: - row = await self.fetchone() - if row is None: + if not self._rows: + await self._async_fetch() + if not self._rows: break - rows.append(row) + self._take_buffered_rows(rows, size) return rows @override diff --git a/pyathena/arrow/result_set.py b/pyathena/arrow/result_set.py index 53939437e..508f644d7 100644 --- a/pyathena/arrow/result_set.py +++ b/pyathena/arrow/result_set.py @@ -216,21 +216,6 @@ def _fetch(self) -> None: processed_rows = list(zip(*columns, strict=False)) self._rows.extend(processed_rows) - @override - def fetchone( - self, - ) -> tuple[Any | None, ...] | dict[Any, Any | None] | None: - if not self._rows: - self._fetch() - if not self._rows: - return None - row = self._rows.popleft() - # As _count_rows() does, inline on the per-row path. - rownumber = self._rownumber - if rownumber is not None: - self._rownumber = rownumber + 1 - return row - def _open_csv(self, location: str) -> Any: """Open the CSV result file for ``pyarrow.csv.read_csv()``. diff --git a/pyathena/result_set.py b/pyathena/result_set.py index a72271fab..3e29b938a 100644 --- a/pyathena/result_set.py +++ b/pyathena/result_set.py @@ -459,14 +459,18 @@ def _get_query_results(self, max_results: int, next_token: str | None = None) -> return cast(dict[str, Any], response) def _fetch(self) -> None: - """Fetch the next page of rows into the result set. + """Fetch the next rows, if any, into the buffer. + + The default result set fetches the next page; it fetches nothing when no + page remains, as after ``close()``. Raises: - ProgrammingError: If there is no next page. + ProgrammingError: If the next page is requested and the query ID is + missing, the query has not succeeded, or the result set is closed. OperationalError: If the request fails. """ if not self._next_token: - raise ProgrammingError("NextToken is none or empty.") + return response = self._get_query_results(self._arraysize, self._next_token) rows, self._next_token = self._parse_result_rows(response) self._process_rows(rows) @@ -490,10 +494,10 @@ def _pre_fetch(self) -> None: def fetchone( self, ) -> tuple[Any | None, ...] | dict[Any, Any | None] | None: - if not self._rows and self._next_token: - self._fetch() if not self._rows: - return None + self._fetch() + if not self._rows: + return None row = self._rows.popleft() # As _count_rows() does, inline on the per-row path. rownumber = self._rownumber @@ -536,9 +540,9 @@ def fetchall( def _take_rows(self, size: int | None) -> list[tuple[Any | None, ...] | dict[Any, Any | None]]: """Take the next rows and count them in ``rownumber``. - The buffered rows are taken in bulk. When the buffer is empty, ``fetchone()`` - takes the next row, so that it fetches the next rows into the buffer as the - result set does. + The buffered rows are taken in bulk, and ``_fetch()`` refills the buffer when + it is empty. A subclass whose ``fetchone()`` does not use the buffer overrides + this method too. Args: size: The maximum number of rows, or None for all remaining rows. @@ -548,13 +552,11 @@ def _take_rows(self, size: int | None) -> list[tuple[Any | None, ...] | dict[Any """ rows: list[tuple[Any | None, ...] | dict[Any, Any | None]] = [] while size is None or len(rows) < size: - if self._rows: - self._take_buffered_rows(rows, size) - else: - row = self.fetchone() - if row is None: + if not self._rows: + self._fetch() + if not self._rows: break - rows.append(row) + self._take_buffered_rows(rows, size) return rows def _take_buffered_rows( diff --git a/pyathena/s3fs/result_set.py b/pyathena/s3fs/result_set.py index 1ae07684a..7569a92f3 100644 --- a/pyathena/s3fs/result_set.py +++ b/pyathena/s3fs/result_set.py @@ -253,26 +253,6 @@ def _fetch(self) -> None: self._rows.append(converted_row) rows_fetched += 1 - @override - def fetchone( - self, - ) -> tuple[Any | None, ...] | dict[Any, Any | None] | None: - """Fetch the next row of the result set. - - Returns: - A tuple representing the next row, or None if no more rows. - """ - if not self._rows: - self._fetch() - if not self._rows: - return None - row = self._rows.popleft() - # As _count_rows() does, inline on the per-row path. - rownumber = self._rownumber - if rownumber is not None: - self._rownumber = rownumber + 1 - return row - @override def close(self) -> None: """Close the result set and release resources.""" diff --git a/tests/pyathena/aio/test_result_set.py b/tests/pyathena/aio/test_result_set.py index 61c27e8f1..90a635dcc 100644 --- a/tests/pyathena/aio/test_result_set.py +++ b/tests/pyathena/aio/test_result_set.py @@ -10,6 +10,7 @@ from pyathena.aio.result_set import AthenaAioDictResultSet, AthenaAioResultSet from pyathena.converter import DefaultTypeConverter +from pyathena.error import OperationalError from pyathena.model import AthenaQueryExecution from pyathena.util import RetryConfig from tests.pyathena.util import get_query_results_responses @@ -20,8 +21,9 @@ async def _create_result_set(result_set_class, pages=(("1", "2"),)): Args: result_set_class: The asyncio result set class to create. - pages: The values of an integer column ``a``, one tuple per page. The first - page is fetched at creation, the others on demand. + pages: The values of an integer column ``a``, one tuple per page, or an + exception that the request for the page raises. The first page is fetched + at creation, the others on demand. Returns: The result set. @@ -90,3 +92,27 @@ async def test_fetchmany_size_not_integer(self): with pytest.raises(TypeError): await result_set.fetchmany(1.5) assert await result_set.fetchall() == [(1,), (2,)] + + async def test_fetchone_across_pages(self): + result_set = await _create_result_set(AthenaAioResultSet, [("1", "2"), ("3",)]) + assert [await result_set.fetchone() for _ in range(4)] == [(1,), (2,), (3,), None] + assert result_set.rownumber == 3 + assert result_set.connection.client.get_query_results.call_count == 2 + + @pytest.mark.parametrize(("method", "args"), [("fetchmany", (3,)), ("fetchall", ())]) + async def test_fetch_counts_rows_taken_before_error(self, method, args): + """The rows taken before a failed page request are counted in rownumber.""" + result_set = await _create_result_set( + AthenaAioResultSet, [("1", "2"), RuntimeError("failed")] + ) + with pytest.raises(OperationalError): + await getattr(result_set, method)(*args) + assert result_set.rownumber == 2 + + async def test_fetch_after_close(self): + result_set = await _create_result_set(AthenaAioResultSet, [("1", "2"), ("3",)]) + result_set.close() + assert await result_set.fetchone() is None + assert await result_set.fetchmany() == [] + assert await result_set.fetchall() == [] + assert result_set.rownumber is None diff --git a/tests/pyathena/arrow/test_result_set.py b/tests/pyathena/arrow/test_result_set.py index 2b8eddf35..46be7dc41 100644 --- a/tests/pyathena/arrow/test_result_set.py +++ b/tests/pyathena/arrow/test_result_set.py @@ -35,6 +35,35 @@ def test_fetch_after_close(self): assert result_set.fetchmany() == [] assert result_set.fetchall() == [] + def test_fetch_across_batches(self): + """The buffer is refilled one record batch at a time. + + No AWS calls; the first page request and the result table are mocked. + """ + with ( + patch.object(AthenaArrowResultSet, "_pre_fetch"), + patch.object(AthenaArrowResultSet, "_create_s3_file_system"), + patch.object( + AthenaArrowResultSet, "_as_arrow", return_value=pa.table({"a": [1, 2, 3]}) + ), + ): + result_set = AthenaArrowResultSet( + connection=MagicMock(), + converter=DefaultArrowTypeConverter(), + query_execution=MagicMock( + state=AthenaQueryExecution.STATE_SUCCEEDED, + output_location="s3://bucket/query.csv", + ), + arraysize=1, + retry_config=RetryConfig(), + ) + assert result_set.fetchmany(2) == [(1,), (2,)] + assert result_set.fetchone() == (3,) + assert result_set.fetchmany() == [] + assert result_set.fetchall() == [] + assert result_set.fetchone() is None + assert result_set.rownumber == 3 + def test_read_csv_timestamps_with_the_same_name(self): """Timestamp columns are converted by position, also with names that repeat. diff --git a/tests/pyathena/s3fs/test_result_set.py b/tests/pyathena/s3fs/test_result_set.py index 4bb1355e9..0c72aa5b7 100644 --- a/tests/pyathena/s3fs/test_result_set.py +++ b/tests/pyathena/s3fs/test_result_set.py @@ -121,6 +121,17 @@ def test_reader_closes_at_end_of_results(self): assert result_set.fetchone() is None result_set.close() + def test_fetch_across_batches(self): + """The buffer is refilled one batch of arraysize rows at a time.""" + result_set, _ = _result_set("null_col,empty_col,value_col\n,,a\n,,b\n,,c\n", None) + assert result_set.fetchmany(2) == [(None, None, "a"), (None, None, "b")] + assert result_set.fetchone() == (None, None, "c") + assert result_set.fetchmany() == [] + assert result_set.fetchall() == [] + assert result_set.fetchone() is None + assert result_set.rownumber == 3 + result_set.close() + def test_reader_failure_closes_file(self): """A failure to read the header row closes the file.""" data = "null_col,empty_col,value_col\n" diff --git a/tests/pyathena/test_result_set.py b/tests/pyathena/test_result_set.py index c0e11238c..087506552 100644 --- a/tests/pyathena/test_result_set.py +++ b/tests/pyathena/test_result_set.py @@ -54,27 +54,6 @@ def _fetching_result_set(*pages): ) -class _FetchoneResultSet(AthenaResultSet): - """A result set whose ``fetchone()`` takes rows from an iterator, not the buffer.""" - - def __init__(self, rows): - """Initialize the result set with its rows, skipping the base initialization. - - Args: - rows: The rows that ``fetchone()`` returns. - """ - self._rows = collections.deque() - self._rownumber = 0 - self._row_iter = iter(rows) - - @override - def fetchone(self): - row = next(self._row_iter, None) - if row is not None: - self._rownumber += 1 - return row - - class _ClosingDeque(collections.deque): """A buffer that closes its result set when it is read, as another thread may.""" @@ -158,12 +137,11 @@ def test_fetch_with_close_during_take(self, method, args, expected): assert getattr(result_set, method)(*args) == expected assert result_set.rownumber is None - def test_fetch_with_fetchone_override(self): - """A subclass whose fetchone() does not use the buffer has its rows fetched.""" - result_set = _FetchoneResultSet([(1,), (2,), (3,)]) - assert result_set.fetchmany(2) == [(1,), (2,)] - assert result_set.fetchall() == [(3,)] + def test_fetchone_across_pages(self): + result_set = _fetching_result_set(("1", "2"), ("3",)) + assert [result_set.fetchone() for _ in range(4)] == [(1,), (2,), (3,), None] assert result_set.rownumber == 3 + assert result_set.connection.client.get_query_results.call_count == 2 def test_fetch_all_rows_skips_column_labels_only_on_first_page(self): """A later page can start with a data row equal to the column labels.