From f78d06c075c98ed91c756185d6500c39ff885df3 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Tue, 29 Sep 2026 01:26:24 +0900 Subject: [PATCH 1/8] Share the pure cursor-base logic between the sync and asyncio cursors The asyncio cursor base copied the result cache search, the execute request preparation, and the result set members of the sync cursor base. Move the parts without I/O to the shared base classes: - BaseCursor._cache_search_limits and _match_previous_query hold the cache window and the per-page match of _find_previous_query_id; both _find_previous_query_id methods keep only the paging and the listing. - BaseCursor._build_execute_request formats the query and builds the StartQueryExecution request for both _execute methods. - WithResultSet now provides the result_set, query_id, arraysize and rownumber properties and the default fetch methods that WithFetch and WithAsyncFetch repeated. Both mixins list WithResultSet before CursorIterator so that these members take precedence. close stays in each mixin because BaseCursor declares it abstract. No-AWS tests pin the cache search's expiry stop, paging, and latest-match behavior for both cursor bases. Refs #880 Co-Authored-By: Claude Opus 5.5 --- pyathena/aio/common.py | 194 ++++++++---------------------- pyathena/common.py | 184 +++++++++++++++++++++------- pyathena/result_set.py | 168 +++++++++++--------------- tests/pyathena/aio/test_cursor.py | 69 ++++++++++- tests/pyathena/test_cursor.py | 56 ++++++++- tests/pyathena/util.py | 20 ++- 6 files changed, 404 insertions(+), 287 deletions(-) diff --git a/pyathena/aio/common.py b/pyathena/aio/common.py index a2beac17..a4849ea0 100644 --- a/pyathena/aio/common.py +++ b/pyathena/aio/common.py @@ -2,10 +2,8 @@ import asyncio import logging -import sys from collections.abc import Awaitable, Callable -from datetime import datetime, timedelta, timezone -from typing import Any, TypeVar, cast +from typing import Any, TypeVar from botocore.exceptions import BotoCoreError, ClientError @@ -15,7 +13,7 @@ from pyathena.glue import GlueMetadataClient from pyathena.model import AthenaDatabase, AthenaQueryExecution, AthenaTableMetadata from pyathena.options import ExecuteOptions -from pyathena.result_set import AthenaResultSet, WithResultSet +from pyathena.result_set import WithResultSet from pyathena.util import _is_throttling_error _logger = logging.getLogger(__name__) @@ -44,6 +42,29 @@ async def _execute( # type: ignore[override] paramstyle: str | None = None, options: ExecuteOptions | None = None, ) -> str: + """Start a query execution, or find a previous one to reuse. + + The individual keyword arguments override the ``options`` field of the + same name unless None. + + Args: + operation: SQL query string. + parameters: Query parameters. + work_group: Athena work group. + s3_staging_dir: S3 location for query results. + cache_size: Number of recent executions to search for a reusable result. + cache_expiration_time: Maximum age of a reusable result in seconds. + result_reuse_enable: Whether to enable Athena result reuse. + result_reuse_minutes: Maximum age of an Athena-reused result in minutes. + paramstyle: Parameter style ('qmark' or 'pyformat'). + options: The execution options. + + Returns: + The query execution ID. + + Raises: + DatabaseError: If the ``StartQueryExecution`` request fails. + """ # The individual keyword arguments are retained for backward compatibility # with external callers that predate ExecuteOptions, mirroring # BaseCursor._execute(). @@ -57,16 +78,7 @@ async def _execute( # type: ignore[override] result_reuse_minutes=result_reuse_minutes, paramstyle=paramstyle, ) - query, execution_parameters = self._prepare_query(operation, parameters, options.paramstyle) - - request = self._build_start_query_execution_request( - query=query, - work_group=options.work_group, - s3_staging_dir=options.s3_staging_dir, - result_reuse_enable=options.result_reuse_enable, - result_reuse_minutes=options.result_reuse_minutes, - execution_parameters=execution_parameters, - ) + query, request = self._build_execute_request(operation, parameters, options) query_id = await self._find_previous_query_id( query, options.work_group, @@ -236,13 +248,23 @@ async def _find_previous_query_id( # type: ignore[override] cache_size: int = 0, cache_expiration_time: int = 0, ) -> str | None: + """Find a previous execution of a query whose result can be reused. + + Searches the work group's recent executions page by page. A failed + search is logged and treated as a cache miss. + + Args: + query: The query string. + work_group: The work group to search, or None for the cursor's. + cache_size: The number of recent executions to search, or 0. + cache_expiration_time: The maximum age of a reused result in + seconds, or 0 for no limit. + + Returns: + The query ID of the latest reusable execution, or None. + """ + cache_size, expiration_time = self._cache_search_limits(cache_size, cache_expiration_time) query_id = None - if cache_size == 0 and cache_expiration_time > 0: - cache_size = sys.maxsize - if cache_expiration_time > 0: - expiration_time = datetime.now(timezone.utc) - timedelta(seconds=cache_expiration_time) - else: - expiration_time = datetime.now(timezone.utc) try: next_token = None while cache_size > 0: @@ -251,32 +273,10 @@ async def _find_previous_query_id( # type: ignore[override] next_token, query_executions = await self._list_query_executions( work_group, next_token=next_token, max_results=max_results ) - for execution in sorted( - ( - e - for e in query_executions - if e.state == AthenaQueryExecution.STATE_SUCCEEDED - and e.statement_type == AthenaQueryExecution.STATEMENT_TYPE_DML - ), - key=lambda e: e.completion_date_time, # type: ignore[arg-type, return-value] - reverse=True, - ): - if ( - cache_expiration_time > 0 - and execution.completion_date_time - and execution.completion_date_time.astimezone(timezone.utc) - < expiration_time - ): - next_token = None - break - if ( - execution.query == query - and execution.database == self._schema_name - and (execution.catalog or "").lower() == (self._catalog_name or "").lower() - ): - query_id = execution.query_id - break - if query_id or next_token is None: + query_id, expired = self._match_previous_query( + query, query_executions, expiration_time + ) + if query_id or expired or next_token is None: break except Exception: _logger.warning("Failed to check the cache. Moving on without cache.", exc_info=True) @@ -616,57 +616,18 @@ async def athena_request( ) -class WithAsyncFetch(AioBaseCursor, CursorIterator, WithResultSet): - """Mixin providing shared fetch, lifecycle, and async protocol for SQL cursors. +class WithAsyncFetch(AioBaseCursor, WithResultSet, CursorIterator): + """Mixin providing async lifecycle and the async protocol for SQL cursors. - Provides properties (``arraysize``, ``result_set``, ``query_id``, - ``rownumber``, ``rowcount``), lifecycle methods (``close``, ``executemany``, - ``cancel``), default sync fetch (for cursors whose result sets load all - data eagerly in ``__init__``), and the async iteration protocol. + Adds ``close``, async ``executemany`` and ``cancel``, async iteration, and + the async context manager protocol to the properties and default sync + fetch methods of ``WithResultSet``. Cursors whose result sets fetch lazily + override the fetch methods with async versions. Subclasses override ``execute()`` and optionally ``__init__`` and format-specific helpers. """ - def __init__(self, **kwargs) -> None: - super().__init__(**kwargs) - self._query_id: str | None = None - self._result_set: AthenaResultSet | None = None - - @property - def arraysize(self) -> int: - return self._arraysize - - @arraysize.setter - def arraysize(self, value: int) -> None: - if value <= 0: - raise ProgrammingError("arraysize must be a positive integer value.") - self._arraysize = value - - @property # type: ignore[override] - def result_set(self) -> AthenaResultSet | None: - return self._result_set - - @result_set.setter - def result_set(self, val) -> None: - self._result_set = val - - @property - def query_id(self) -> str | None: - return self._query_id - - @query_id.setter - def query_id(self, val) -> None: - self._query_id = val - - @property - def rownumber(self) -> int | None: - return self.result_set.rownumber if self.result_set else None - - @property - def rowcount(self) -> int: - return self.result_set.rowcount if self.result_set else self._rowcount - def close(self) -> None: """Close the cursor and release associated resources.""" self._rowcount = -1 @@ -720,57 +681,6 @@ async def cancel(self) -> None: raise ProgrammingError("QueryExecutionId is none or empty.") await self._cancel(self.query_id) - 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. - - Raises: - ProgrammingError: If no result set is available. - """ - if not self.has_result_set: - raise ProgrammingError("No result set.") - result_set = cast(AthenaResultSet, self.result_set) - return result_set.fetchone() - - def fetchmany( - self, size: int | None = None - ) -> list[tuple[Any | None, ...] | dict[Any, Any | None]]: - """Fetch multiple rows from the result set. - - Args: - size: Maximum number of rows to fetch. Defaults to arraysize. - - Returns: - List of tuples representing the fetched rows. - - Raises: - ProgrammingError: If no result set is available. - """ - if not self.has_result_set: - raise ProgrammingError("No result set.") - result_set = cast(AthenaResultSet, self.result_set) - return result_set.fetchmany(size) - - def fetchall( - self, - ) -> list[tuple[Any | None, ...] | dict[Any, Any | None]]: - """Fetch all remaining rows from the result set. - - Returns: - List of tuples representing all remaining rows. - - Raises: - ProgrammingError: If no result set is available. - """ - if not self.has_result_set: - raise ProgrammingError("No result set.") - result_set = cast(AthenaResultSet, self.result_set) - return result_set.fetchall() - def __aiter__(self): return self diff --git a/pyathena/common.py b/pyathena/common.py index b18868c3..15a9553d 100644 --- a/pyathena/common.py +++ b/pyathena/common.py @@ -914,6 +914,76 @@ def _poll(self, query_id: str) -> AthenaQueryExecution | AthenaCalculationExecut raise e return query_execution + def _cache_search_limits( + self, cache_size: int, cache_expiration_time: int + ) -> tuple[int, datetime | None]: + """Resolve how far the result cache search looks back. No I/O. + + Args: + cache_size: The number of recent executions to search, or 0. + cache_expiration_time: The maximum age of a reused result in + seconds, or 0 for no limit. + + Returns: + The number of executions to search, unbounded when only + ``cache_expiration_time`` is set, and the oldest completion time + to accept, or None for no limit. + """ + if cache_size == 0 and cache_expiration_time > 0: + cache_size = sys.maxsize + if cache_expiration_time > 0: + expiration_time = datetime.now(timezone.utc) - timedelta(seconds=cache_expiration_time) + return cache_size, expiration_time + return cache_size, None + + def _match_previous_query( + self, + query: str, + query_executions: list[AthenaQueryExecution], + expiration_time: datetime | None, + ) -> tuple[str | None, bool]: + """Find the latest reusable execution of a query in one page. No I/O. + + Reusable executions are succeeded DML queries with the same query + string, schema, and catalog (case-insensitive). Executions are checked + from the latest completion; the check stops at the first one completed + before ``expiration_time``. + + Args: + query: The query string. + query_executions: One page of the work group's executions. + expiration_time: The oldest completion time to accept, or None for + no limit. + + Returns: + The matching query ID or None, and whether the check reached an + expired execution. + """ + for execution in sorted( + ( + e + for e in query_executions + if e.state == AthenaQueryExecution.STATE_SUCCEEDED + and e.statement_type == AthenaQueryExecution.STATEMENT_TYPE_DML + ), + # https://github.com/python/mypy/issues/9656 + key=lambda e: e.completion_date_time, # type: ignore[arg-type, return-value] + reverse=True, + ): + if ( + expiration_time + and execution.completion_date_time + and execution.completion_date_time.astimezone(timezone.utc) < expiration_time + ): + return None, True + if ( + execution.query == query + and execution.database == self._schema_name + and (execution.catalog or "").lower() == (self._catalog_name or "").lower() + ): + return execution.query_id, False + return None, False + def _find_previous_query_id( self, query: str, @@ -921,13 +991,23 @@ def _find_previous_query_id( cache_size: int = 0, cache_expiration_time: int = 0, ) -> str | None: + """Find a previous execution of a query whose result can be reused. + + Searches the work group's recent executions page by page. A failed + search is logged and treated as a cache miss. + + Args: + query: The query string. + work_group: The work group to search, or None for the cursor's. + cache_size: The number of recent executions to search, or 0. + cache_expiration_time: The maximum age of a reused result in + seconds, or 0 for no limit. + + Returns: + The query ID of the latest reusable execution, or None. + """ + cache_size, expiration_time = self._cache_search_limits(cache_size, cache_expiration_time) query_id = None - if cache_size == 0 and cache_expiration_time > 0: - cache_size = sys.maxsize - if cache_expiration_time > 0: - expiration_time = datetime.now(timezone.utc) - timedelta(seconds=cache_expiration_time) - else: - expiration_time = datetime.now(timezone.utc) try: next_token = None while cache_size > 0: @@ -936,33 +1016,10 @@ def _find_previous_query_id( next_token, query_executions = self._list_query_executions( work_group, next_token=next_token, max_results=max_results ) - for execution in sorted( - ( - e - for e in query_executions - if e.state == AthenaQueryExecution.STATE_SUCCEEDED - and e.statement_type == AthenaQueryExecution.STATEMENT_TYPE_DML - ), - # https://github.com/python/mypy/issues/9656 - key=lambda e: e.completion_date_time, # type: ignore[arg-type, return-value] - reverse=True, - ): - if ( - cache_expiration_time > 0 - and execution.completion_date_time - and execution.completion_date_time.astimezone(timezone.utc) - < expiration_time - ): - next_token = None - break - if ( - execution.query == query - and execution.database == self._schema_name - and (execution.catalog or "").lower() == (self._catalog_name or "").lower() - ): - query_id = execution.query_id - break - if query_id or next_token is None: + query_id, expired = self._match_previous_query( + query, query_executions, expiration_time + ) + if query_id or expired or next_token is None: break except Exception: _logger.warning("Failed to check the cache. Moving on without cache.", exc_info=True) @@ -1031,6 +1088,33 @@ def _call_on_start_query_execution(self, query_id: str, options: ExecuteOptions) if options.on_start_query_execution: options.on_start_query_execution(query_id) + def _build_execute_request( + self, + operation: str, + parameters: dict[str, Any] | list[str] | None, + options: ExecuteOptions, + ) -> tuple[str, dict[str, Any]]: + """Format a query and build its ``StartQueryExecution`` request. No I/O. + + Args: + operation: SQL query string. + parameters: Query parameters. + options: The resolved execution options. + + Returns: + Tuple of (formatted_query, request). + """ + query, execution_parameters = self._prepare_query(operation, parameters, options.paramstyle) + request = self._build_start_query_execution_request( + query=query, + work_group=options.work_group, + s3_staging_dir=options.s3_staging_dir, + result_reuse_enable=options.result_reuse_enable, + result_reuse_minutes=options.result_reuse_minutes, + execution_parameters=execution_parameters, + ) + return query, request + def _execute( self, operation: str, @@ -1044,6 +1128,29 @@ def _execute( paramstyle: str | None = None, options: ExecuteOptions | None = None, ) -> str: + """Start a query execution, or find a previous one to reuse. + + The individual keyword arguments override the ``options`` field of the + same name unless None. + + Args: + operation: SQL query string. + parameters: Query parameters. + work_group: Athena work group. + s3_staging_dir: S3 location for query results. + cache_size: Number of recent executions to search for a reusable result. + cache_expiration_time: Maximum age of a reusable result in seconds. + result_reuse_enable: Whether to enable Athena result reuse. + result_reuse_minutes: Maximum age of an Athena-reused result in minutes. + paramstyle: Parameter style ('qmark' or 'pyformat'). + options: The execution options. + + Returns: + The query execution ID. + + Raises: + DatabaseError: If the ``StartQueryExecution`` request fails. + """ # The individual keyword arguments are retained for backward compatibility # with external callers that predate ExecuteOptions (e.g. dbt-athena <= 1.10.x # calls _execute() with work_group/s3_staging_dir/cache_* keywords). @@ -1057,16 +1164,7 @@ def _execute( result_reuse_minutes=result_reuse_minutes, paramstyle=paramstyle, ) - query, execution_parameters = self._prepare_query(operation, parameters, options.paramstyle) - - request = self._build_start_query_execution_request( - query=query, - work_group=options.work_group, - s3_staging_dir=options.s3_staging_dir, - result_reuse_enable=options.result_reuse_enable, - result_reuse_minutes=options.result_reuse_minutes, - execution_parameters=execution_parameters, - ) + query, request = self._build_execute_request(operation, parameters, options) query_id = self._find_previous_query_id( query, options.work_group, diff --git a/pyathena/result_set.py b/pyathena/result_set.py index d75326c3..76162de3 100644 --- a/pyathena/result_set.py +++ b/pyathena/result_set.py @@ -2,7 +2,6 @@ import collections import logging -from abc import abstractmethod from datetime import datetime from typing import ( TYPE_CHECKING, @@ -796,8 +795,17 @@ def _get_rows( class WithResultSet: + """Mixin providing a SQL cursor's result set, its properties, and default fetch. + + Shared by the sync (``WithFetch``) and asyncio (``WithAsyncFetch``) cursors. + The default fetch methods suit result sets that load all data eagerly in + ``__init__``. + """ + def __init__(self): super().__init__() + self._query_id: str | None = None + self._result_set: AthenaResultSet | None = None def _reset_state(self) -> None: self._rowcount = -1 @@ -807,14 +815,12 @@ def _reset_state(self) -> None: self.result_set = None @property - @abstractmethod def result_set(self) -> AthenaResultSet | None: - raise NotImplementedError # pragma: no cover + return self._result_set @result_set.setter - @abstractmethod def result_set(self, val: AthenaResultSet | None) -> None: - raise NotImplementedError # pragma: no cover + self._result_set = val @property def has_result_set(self) -> bool: @@ -841,14 +847,12 @@ def catalog(self) -> str | None: return self.result_set.catalog @property - @abstractmethod def query_id(self) -> str | None: - raise NotImplementedError # pragma: no cover + return self._query_id @query_id.setter - @abstractmethod def query_id(self, val: str | None) -> None: - raise NotImplementedError # pragma: no cover + self._query_id = val @property def query(self) -> str | None: @@ -1044,24 +1048,6 @@ def rowcount(self) -> int: """ return self.result_set.rowcount if self.result_set else self._rowcount - -class WithFetch(BaseCursor, CursorIterator, WithResultSet): - """Mixin providing shared properties, fetch, lifecycle, and sync iteration for SQL cursors. - - Provides properties (``arraysize``, ``result_set``, ``query_id``, - ``rownumber``, ``rowcount``), lifecycle methods (``close``, ``executemany``, - ``cancel``), default sync fetch (for cursors whose result sets load all - data eagerly in ``__init__``), and the sync iteration protocol. - - Subclasses override ``execute()`` and optionally ``__init__`` and - format-specific helpers. - """ - - def __init__(self, **kwargs) -> None: - super().__init__(**kwargs) - self._query_id: str | None = None - self._result_set: AthenaResultSet | None = None - @property def arraysize(self) -> int: return self._arraysize @@ -1073,28 +1059,71 @@ def arraysize(self, value: int) -> None: self._arraysize = value @property - def result_set(self) -> AthenaResultSet | None: - return self._result_set + def rownumber(self) -> int | None: + return self.result_set.rownumber if self.result_set else None - @result_set.setter - def result_set(self, val) -> None: - self._result_set = val + def fetchone( + self, + ) -> tuple[Any | None, ...] | dict[Any, Any | None] | None: + """Fetch the next row of the result set. - @property - def query_id(self) -> str | None: - return self._query_id + Returns: + A tuple representing the next row, or None if no more rows. - @query_id.setter - def query_id(self, val) -> None: - self._query_id = val + Raises: + ProgrammingError: If no result set is available. + """ + if not self.has_result_set: + raise ProgrammingError("No result set.") + result_set = cast(AthenaResultSet, self.result_set) + return result_set.fetchone() - @property - def rownumber(self) -> int | None: - return self.result_set.rownumber if self.result_set else None + def fetchmany( + self, size: int | None = None + ) -> list[tuple[Any | None, ...] | dict[Any, Any | None]]: + """Fetch multiple rows from the result set. - @property - def rowcount(self) -> int: - return self.result_set.rowcount if self.result_set else self._rowcount + Args: + size: Maximum number of rows to fetch. Defaults to arraysize. + + Returns: + List of tuples representing the fetched rows. + + Raises: + ProgrammingError: If no result set is available. + """ + if not self.has_result_set: + raise ProgrammingError("No result set.") + result_set = cast(AthenaResultSet, self.result_set) + return result_set.fetchmany(size) + + def fetchall( + self, + ) -> list[tuple[Any | None, ...] | dict[Any, Any | None]]: + """Fetch all remaining rows from the result set. + + Returns: + List of tuples representing all remaining rows. + + Raises: + ProgrammingError: If no result set is available. + """ + if not self.has_result_set: + raise ProgrammingError("No result set.") + result_set = cast(AthenaResultSet, self.result_set) + return result_set.fetchall() + + +class WithFetch(BaseCursor, WithResultSet, CursorIterator): + """Mixin providing sync lifecycle and iteration for SQL cursors. + + Adds ``close``, ``executemany``, and ``cancel`` to the properties and + default fetch methods of ``WithResultSet``, and the sync iteration + protocol through ``CursorIterator``. + + Subclasses override ``execute()`` and optionally ``__init__`` and + format-specific helpers. + """ def close(self) -> None: """Close the cursor and release associated resources.""" @@ -1148,54 +1177,3 @@ def cancel(self) -> None: if not self.query_id: raise ProgrammingError("QueryExecutionId is none or empty.") self._cancel(self.query_id) - - 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. - - Raises: - ProgrammingError: If no result set is available. - """ - if not self.has_result_set: - raise ProgrammingError("No result set.") - result_set = cast(AthenaResultSet, self.result_set) - return result_set.fetchone() - - def fetchmany( - self, size: int | None = None - ) -> list[tuple[Any | None, ...] | dict[Any, Any | None]]: - """Fetch multiple rows from the result set. - - Args: - size: Maximum number of rows to fetch. Defaults to arraysize. - - Returns: - List of tuples representing the fetched rows. - - Raises: - ProgrammingError: If no result set is available. - """ - if not self.has_result_set: - raise ProgrammingError("No result set.") - result_set = cast(AthenaResultSet, self.result_set) - return result_set.fetchmany(size) - - def fetchall( - self, - ) -> list[tuple[Any | None, ...] | dict[Any, Any | None]]: - """Fetch all remaining rows from the result set. - - Returns: - List of tuples representing all remaining rows. - - Raises: - ProgrammingError: If no result set is available. - """ - if not self.has_result_set: - raise ProgrammingError("No result set.") - result_set = cast(AthenaResultSet, self.result_set) - return result_set.fetchall() diff --git a/tests/pyathena/aio/test_cursor.py b/tests/pyathena/aio/test_cursor.py index a1259eff..da587fbd 100644 --- a/tests/pyathena/aio/test_cursor.py +++ b/tests/pyathena/aio/test_cursor.py @@ -1,8 +1,8 @@ import asyncio import re import threading -from datetime import datetime, timezone -from unittest.mock import AsyncMock, MagicMock, patch +from datetime import datetime, timedelta, timezone +from unittest.mock import AsyncMock, MagicMock, call, patch import pytest @@ -15,7 +15,7 @@ from pyathena.util import RetryConfig from tests import ENV from tests.pyathena.aio.conftest import _aio_connect -from tests.pyathena.util import throttle_metadata_api +from tests.pyathena.util import succeeded_query_execution, throttle_metadata_api class TestAioCursor: @@ -233,6 +233,69 @@ def execution(catalog): == "query_id_awsdatacatalog" ) + async def test_cache_search_stops_at_expired_execution(self): + query = "SELECT * FROM one_row" + now = datetime.now(timezone.utc) + cursor = AioCursor.__new__(AioCursor) + cursor._schema_name = "this_schema" + cursor._catalog_name = None + page = [ + succeeded_query_execution("expired", query, now - timedelta(hours=2)), + succeeded_query_execution("other", "SELECT 1", now), + ] + + with patch.object( + AioCursor, + "_list_query_executions", + new_callable=AsyncMock, + return_value=("next_token", page), + ) as list_mock: + # Without cache_size, cache_expiration_time alone bounds the search. + assert ( + await cursor._find_previous_query_id(query, None, cache_expiration_time=3600) + is None + ) + list_mock.assert_awaited_once_with(None, next_token=None, max_results=50) + + async def test_cache_search_reads_pages_up_to_cache_size(self): + query = "SELECT * FROM one_row" + now = datetime.now(timezone.utc) + cursor = AioCursor.__new__(AioCursor) + cursor._schema_name = "this_schema" + cursor._catalog_name = None + pages = [ + ("next_token", [succeeded_query_execution("other", "SELECT 1", now)]), + ("last_token", [succeeded_query_execution("match", query, now)]), + ] + + with patch.object( + AioCursor, "_list_query_executions", new_callable=AsyncMock, side_effect=pages + ) as list_mock: + assert await cursor._find_previous_query_id(query, "wg", cache_size=60) == "match" + assert list_mock.await_args_list == [ + call("wg", next_token=None, max_results=50), + call("wg", next_token="next_token", max_results=10), + ] + + async def test_cache_search_prefers_latest_execution(self): + query = "SELECT * FROM one_row" + now = datetime.now(timezone.utc) + cursor = AioCursor.__new__(AioCursor) + cursor._schema_name = "this_schema" + cursor._catalog_name = None + page = [ + succeeded_query_execution("older", query, now - timedelta(minutes=1)), + succeeded_query_execution("latest", query, now), + ] + + with patch.object( + AioCursor, + "_list_query_executions", + new_callable=AsyncMock, + return_value=(None, page), + ): + assert await cursor._find_previous_query_id(query, None, cache_size=10) == "latest" + async def test_no_result_set_raises(self, aio_cursor): with pytest.raises(ProgrammingError): await aio_cursor.fetchone() diff --git a/tests/pyathena/test_cursor.py b/tests/pyathena/test_cursor.py index aa1832b3..8d7f093f 100644 --- a/tests/pyathena/test_cursor.py +++ b/tests/pyathena/test_cursor.py @@ -9,10 +9,10 @@ import uuid from concurrent import futures from concurrent.futures.thread import ThreadPoolExecutor -from datetime import date, datetime, timezone +from datetime import date, datetime, timedelta, timezone from decimal import Decimal from random import randint -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock, call, patch import pytest from botocore.exceptions import ClientError @@ -36,7 +36,7 @@ from pyathena.util import RetryConfig from tests import ENV from tests.pyathena.conftest import connect -from tests.pyathena.util import throttle_metadata_api, unreachable_glue +from tests.pyathena.util import succeeded_query_execution, throttle_metadata_api, unreachable_glue _logger = logging.getLogger(__name__) @@ -277,6 +277,56 @@ def execution(catalog): == "query_id_awsdatacatalog" ) + def test_cache_search_stops_at_expired_execution(self): + query = "SELECT * FROM one_row" + now = datetime.now(timezone.utc) + cursor = Cursor.__new__(Cursor) + cursor._schema_name = "this_schema" + cursor._catalog_name = None + page = [ + succeeded_query_execution("expired", query, now - timedelta(hours=2)), + succeeded_query_execution("other", "SELECT 1", now), + ] + + with patch.object( + Cursor, "_list_query_executions", return_value=("next_token", page) + ) as list_mock: + # Without cache_size, cache_expiration_time alone bounds the search. + assert cursor._find_previous_query_id(query, None, cache_expiration_time=3600) is None + list_mock.assert_called_once_with(None, next_token=None, max_results=50) + + def test_cache_search_reads_pages_up_to_cache_size(self): + query = "SELECT * FROM one_row" + now = datetime.now(timezone.utc) + cursor = Cursor.__new__(Cursor) + cursor._schema_name = "this_schema" + cursor._catalog_name = None + pages = [ + ("next_token", [succeeded_query_execution("other", "SELECT 1", now)]), + ("last_token", [succeeded_query_execution("match", query, now)]), + ] + + with patch.object(Cursor, "_list_query_executions", side_effect=pages) as list_mock: + assert cursor._find_previous_query_id(query, "wg", cache_size=60) == "match" + assert list_mock.call_args_list == [ + call("wg", next_token=None, max_results=50), + call("wg", next_token="next_token", max_results=10), + ] + + def test_cache_search_prefers_latest_execution(self): + query = "SELECT * FROM one_row" + now = datetime.now(timezone.utc) + cursor = Cursor.__new__(Cursor) + cursor._schema_name = "this_schema" + cursor._catalog_name = None + page = [ + succeeded_query_execution("older", query, now - timedelta(minutes=1)), + succeeded_query_execution("latest", query, now), + ] + + with patch.object(Cursor, "_list_query_executions", return_value=(None, page)): + assert cursor._find_previous_query_id(query, None, cache_size=10) == "latest" + @pytest.mark.parametrize( "cursor", [{"work_group": ENV.work_group, "result_reuse_enable": True, "result_reuse_minutes": 5}], diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index f5ac673d..a0f52268 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -14,7 +14,7 @@ from sqlalchemy import types from pyathena.glue import GlueMetadataClient -from pyathena.model import AthenaCalculationExecutionStatus +from pyathena.model import AthenaCalculationExecutionStatus, AthenaQueryExecution _queries = Environment( loader=FileSystemLoader(Path(__file__).parents[1].resolve() / "resources" / "queries") @@ -51,6 +51,24 @@ def fail(**kwargs): return calls +def succeeded_query_execution(query_id, query, completion_date_time, schema="this_schema"): + """A succeeded DML query execution, as the result cache search lists it.""" + return AthenaQueryExecution( + { + "QueryExecution": { + "QueryExecutionId": query_id, + "Query": query, + "StatementType": AthenaQueryExecution.STATEMENT_TYPE_DML, + "QueryExecutionContext": {"Database": schema}, + "Status": { + "State": AthenaQueryExecution.STATE_SUCCEEDED, + "CompletionDateTime": completion_date_time, + }, + } + } + ) + + def unreachable_glue(connection): """A Glue client for the connection that sends requests to a closed proxy port.""" return GlueMetadataClient( From 84a46f52c89236e80badcd6eef0f7675731ae6e9 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Tue, 29 Sep 2026 01:29:05 +0900 Subject: [PATCH 2/8] Fail rather than page on in the cache expiry tests Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/aio/test_cursor.py | 5 ++++- tests/pyathena/test_cursor.py | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/pyathena/aio/test_cursor.py b/tests/pyathena/aio/test_cursor.py index da587fbd..7df4fe51 100644 --- a/tests/pyathena/aio/test_cursor.py +++ b/tests/pyathena/aio/test_cursor.py @@ -244,11 +244,14 @@ async def test_cache_search_stops_at_expired_execution(self): succeeded_query_execution("other", "SELECT 1", now), ] + # A page read after the expired execution would return this match. + next_page = (None, [succeeded_query_execution("next_page", query, now)]) + with patch.object( AioCursor, "_list_query_executions", new_callable=AsyncMock, - return_value=("next_token", page), + side_effect=[("next_token", page), next_page], ) as list_mock: # Without cache_size, cache_expiration_time alone bounds the search. assert ( diff --git a/tests/pyathena/test_cursor.py b/tests/pyathena/test_cursor.py index 8d7f093f..735e7628 100644 --- a/tests/pyathena/test_cursor.py +++ b/tests/pyathena/test_cursor.py @@ -288,8 +288,11 @@ def test_cache_search_stops_at_expired_execution(self): succeeded_query_execution("other", "SELECT 1", now), ] + # A page read after the expired execution would return this match. + next_page = (None, [succeeded_query_execution("next_page", query, now)]) + with patch.object( - Cursor, "_list_query_executions", return_value=("next_token", page) + Cursor, "_list_query_executions", side_effect=[("next_token", page), next_page] ) as list_mock: # Without cache_size, cache_expiration_time alone bounds the search. assert cursor._find_previous_query_id(query, None, cache_expiration_time=3600) is None From a4a5d61099f10182071e15e7cab1081e8966d270 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Tue, 29 Sep 2026 01:31:46 +0900 Subject: [PATCH 3/8] Correct the fetch and error descriptions in the cursor base docstrings Co-Authored-By: Claude Opus 5.5 --- pyathena/aio/common.py | 7 ++++--- pyathena/common.py | 4 ++++ pyathena/result_set.py | 4 ++-- 3 files changed, 10 insertions(+), 5 deletions(-) diff --git a/pyathena/aio/common.py b/pyathena/aio/common.py index a4849ea0..bea5c44c 100644 --- a/pyathena/aio/common.py +++ b/pyathena/aio/common.py @@ -63,6 +63,7 @@ async def _execute( # type: ignore[override] The query execution ID. Raises: + ProgrammingError: If the formatter rejects the query or its parameters. DatabaseError: If the ``StartQueryExecution`` request fails. """ # The individual keyword arguments are retained for backward compatibility @@ -620,9 +621,9 @@ class WithAsyncFetch(AioBaseCursor, WithResultSet, CursorIterator): """Mixin providing async lifecycle and the async protocol for SQL cursors. Adds ``close``, async ``executemany`` and ``cancel``, async iteration, and - the async context manager protocol to the properties and default sync - fetch methods of ``WithResultSet``. Cursors whose result sets fetch lazily - override the fetch methods with async versions. + the async context manager protocol to the properties and sync fetch + methods of ``WithResultSet``. Subclasses override the fetch methods with + async versions. Subclasses override ``execute()`` and optionally ``__init__`` and format-specific helpers. diff --git a/pyathena/common.py b/pyathena/common.py index 15a9553d..7ce367c6 100644 --- a/pyathena/common.py +++ b/pyathena/common.py @@ -1103,6 +1103,9 @@ def _build_execute_request( Returns: Tuple of (formatted_query, request). + + Raises: + ProgrammingError: If the formatter rejects the query or its parameters. """ query, execution_parameters = self._prepare_query(operation, parameters, options.paramstyle) request = self._build_start_query_execution_request( @@ -1149,6 +1152,7 @@ def _execute( The query execution ID. Raises: + ProgrammingError: If the formatter rejects the query or its parameters. DatabaseError: If the ``StartQueryExecution`` request fails. """ # The individual keyword arguments are retained for backward compatibility diff --git a/pyathena/result_set.py b/pyathena/result_set.py index 76162de3..5264cb00 100644 --- a/pyathena/result_set.py +++ b/pyathena/result_set.py @@ -798,8 +798,8 @@ class WithResultSet: """Mixin providing a SQL cursor's result set, its properties, and default fetch. Shared by the sync (``WithFetch``) and asyncio (``WithAsyncFetch``) cursors. - The default fetch methods suit result sets that load all data eagerly in - ``__init__``. + The fetch methods call the result set synchronously; the asyncio cursors + override them with async versions. """ def __init__(self): From ae77fbcb6d43ee529473e057d76694405786359d Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Tue, 29 Sep 2026 01:53:44 +0900 Subject: [PATCH 4/8] Catch a failing expiry check in the cache search tests Co-Authored-By: Claude Opus 5.5 --- pyathena/result_set.py | 1 + tests/pyathena/aio/test_cursor.py | 4 +++- tests/pyathena/test_cursor.py | 4 +++- tests/pyathena/util.py | 12 +++++++++++- 4 files changed, 18 insertions(+), 3 deletions(-) diff --git a/pyathena/result_set.py b/pyathena/result_set.py index 5264cb00..bfca7c58 100644 --- a/pyathena/result_set.py +++ b/pyathena/result_set.py @@ -803,6 +803,7 @@ class WithResultSet: """ def __init__(self): + """Initialize the mixin with no query ID and no result set.""" super().__init__() self._query_id: str | None = None self._result_set: AthenaResultSet | None = None diff --git a/tests/pyathena/aio/test_cursor.py b/tests/pyathena/aio/test_cursor.py index 7df4fe51..7bdab825 100644 --- a/tests/pyathena/aio/test_cursor.py +++ b/tests/pyathena/aio/test_cursor.py @@ -233,7 +233,7 @@ def execution(catalog): == "query_id_awsdatacatalog" ) - async def test_cache_search_stops_at_expired_execution(self): + async def test_cache_search_stops_at_expired_execution(self, caplog): query = "SELECT * FROM one_row" now = datetime.now(timezone.utc) cursor = AioCursor.__new__(AioCursor) @@ -259,6 +259,8 @@ async def test_cache_search_stops_at_expired_execution(self): is None ) list_mock.assert_awaited_once_with(None, next_token=None, max_results=50) + # A failed search also returns None; the expiry must stop it without an error. + assert "Failed to check the cache" not in caplog.text async def test_cache_search_reads_pages_up_to_cache_size(self): query = "SELECT * FROM one_row" diff --git a/tests/pyathena/test_cursor.py b/tests/pyathena/test_cursor.py index 735e7628..cec001d6 100644 --- a/tests/pyathena/test_cursor.py +++ b/tests/pyathena/test_cursor.py @@ -277,7 +277,7 @@ def execution(catalog): == "query_id_awsdatacatalog" ) - def test_cache_search_stops_at_expired_execution(self): + def test_cache_search_stops_at_expired_execution(self, caplog): query = "SELECT * FROM one_row" now = datetime.now(timezone.utc) cursor = Cursor.__new__(Cursor) @@ -297,6 +297,8 @@ def test_cache_search_stops_at_expired_execution(self): # Without cache_size, cache_expiration_time alone bounds the search. assert cursor._find_previous_query_id(query, None, cache_expiration_time=3600) is None list_mock.assert_called_once_with(None, next_token=None, max_results=50) + # A failed search also returns None; the expiry must stop it without an error. + assert "Failed to check the cache" not in caplog.text def test_cache_search_reads_pages_up_to_cache_size(self): query = "SELECT * FROM one_row" diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index a0f52268..bfe4dd3c 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -52,7 +52,17 @@ def fail(**kwargs): def succeeded_query_execution(query_id, query, completion_date_time, schema="this_schema"): - """A succeeded DML query execution, as the result cache search lists it.""" + """Build a succeeded DML query execution, as the result cache search lists it. + + Args: + query_id: The query execution ID. + query: The query string. + completion_date_time: The completion time. + schema: The database the query ran against. + + Returns: + The query execution. + """ return AthenaQueryExecution( { "QueryExecution": { From 7d4562a1e539797044206a70882372eca3cee01f Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Tue, 29 Sep 2026 23:51:45 +0900 Subject: [PATCH 5/8] Fold WithFetch into WithResultSet After the shared members moved to WithResultSet, WithFetch held only the sync close, executemany, and cancel. WithResultSet now subclasses BaseCursor and CursorIterator and holds them, and the sync SQL cursors subclass it directly. WithAsyncFetch subclasses AioBaseCursor and WithResultSet and overrides executemany and cancel with async versions. Because WithResultSet now precedes BaseCursor in the MRO, its close implements the abstract BaseCursor.close, so the copy in each mixin is gone. Breaking change: pyathena.result_set.WithFetch is removed, and a class that lists BaseCursor before WithResultSet in its bases can no longer be created. Refs #880 Co-Authored-By: Claude Opus 5.5 --- pyathena/aio/common.py | 21 +++++++-------------- pyathena/arrow/cursor.py | 4 ++-- pyathena/cursor.py | 4 ++-- pyathena/pandas/cursor.py | 4 ++-- pyathena/polars/cursor.py | 4 ++-- pyathena/result_set.py | 33 +++++++++++++-------------------- pyathena/s3fs/cursor.py | 4 ++-- 7 files changed, 30 insertions(+), 44 deletions(-) diff --git a/pyathena/aio/common.py b/pyathena/aio/common.py index bea5c44c..8510f155 100644 --- a/pyathena/aio/common.py +++ b/pyathena/aio/common.py @@ -8,7 +8,7 @@ from botocore.exceptions import BotoCoreError, ClientError from pyathena.aio.util import async_retry_api_call -from pyathena.common import BaseCursor, CursorIterator +from pyathena.common import BaseCursor from pyathena.error import DatabaseError, OperationalError, ProgrammingError from pyathena.glue import GlueMetadataClient from pyathena.model import AthenaDatabase, AthenaQueryExecution, AthenaTableMetadata @@ -617,24 +617,17 @@ async def athena_request( ) -class WithAsyncFetch(AioBaseCursor, WithResultSet, CursorIterator): - """Mixin providing async lifecycle and the async protocol for SQL cursors. +class WithAsyncFetch(AioBaseCursor, WithResultSet): + """Base class of the asyncio SQL cursors. - Adds ``close``, async ``executemany`` and ``cancel``, async iteration, and - the async context manager protocol to the properties and sync fetch - methods of ``WithResultSet``. Subclasses override the fetch methods with - async versions. + Overrides ``executemany`` and ``cancel`` of ``WithResultSet`` with async + versions and adds async iteration and the async context manager protocol. + Subclasses override the fetch methods with async versions. Subclasses override ``execute()`` and optionally ``__init__`` and format-specific helpers. """ - def close(self) -> None: - """Close the cursor and release associated resources.""" - self._rowcount = -1 - if self.result_set and not self.result_set.is_closed: - self.result_set.close() - async def executemany( # type: ignore[override] self, operation: str, @@ -672,7 +665,7 @@ async def executemany( # type: ignore[override] self._reset_state() self._rowcount = rowcount - async def cancel(self) -> None: + async def cancel(self) -> None: # type: ignore[override] """Cancel the currently executing query. Raises: diff --git a/pyathena/arrow/cursor.py b/pyathena/arrow/cursor.py index 5ad7cd60..6b2dabda 100644 --- a/pyathena/arrow/cursor.py +++ b/pyathena/arrow/cursor.py @@ -13,7 +13,7 @@ from pyathena.error import OperationalError, ProgrammingError from pyathena.model import AthenaQueryExecution from pyathena.options import ExecuteOptions -from pyathena.result_set import WithFetch +from pyathena.result_set import WithResultSet if TYPE_CHECKING: import polars as pl @@ -22,7 +22,7 @@ _logger = logging.getLogger(__name__) -class ArrowCursor(WithFetch): +class ArrowCursor(WithResultSet): """Cursor for handling Apache Arrow Table results from Athena queries. This cursor returns query results as Apache Arrow Tables, which provide diff --git a/pyathena/cursor.py b/pyathena/cursor.py index eb302b30..cc365ce0 100644 --- a/pyathena/cursor.py +++ b/pyathena/cursor.py @@ -8,12 +8,12 @@ from pyathena.error import OperationalError, ProgrammingError from pyathena.model import AthenaQueryExecution from pyathena.options import ExecuteOptions -from pyathena.result_set import AthenaDictResultSet, AthenaResultSet, WithFetch +from pyathena.result_set import AthenaDictResultSet, AthenaResultSet, WithResultSet _logger = logging.getLogger(__name__) -class Cursor(WithFetch): +class Cursor(WithResultSet): """A DB API 2.0 compliant cursor for executing SQL queries on Amazon Athena. The Cursor class provides methods for executing SQL queries against Amazon Athena diff --git a/pyathena/pandas/cursor.py b/pyathena/pandas/cursor.py index 883172c0..38b66c98 100644 --- a/pyathena/pandas/cursor.py +++ b/pyathena/pandas/cursor.py @@ -18,7 +18,7 @@ DefaultPandasUnloadTypeConverter, ) from pyathena.pandas.result_set import AthenaPandasResultSet, PandasDataFrameIterator -from pyathena.result_set import WithFetch +from pyathena.result_set import WithResultSet if TYPE_CHECKING: from pandas import DataFrame @@ -26,7 +26,7 @@ _logger = logging.getLogger(__name__) -class PandasCursor(WithFetch): +class PandasCursor(WithResultSet): """Cursor for handling pandas DataFrame results from Athena queries. This cursor returns query results as pandas DataFrames with memory-efficient diff --git a/pyathena/polars/cursor.py b/pyathena/polars/cursor.py index a92ed580..4bd8c78e 100644 --- a/pyathena/polars/cursor.py +++ b/pyathena/polars/cursor.py @@ -18,7 +18,7 @@ DefaultPolarsUnloadTypeConverter, ) from pyathena.polars.result_set import AthenaPolarsResultSet -from pyathena.result_set import WithFetch +from pyathena.result_set import WithResultSet if TYPE_CHECKING: import polars as pl @@ -27,7 +27,7 @@ _logger = logging.getLogger(__name__) -class PolarsCursor(WithFetch): +class PolarsCursor(WithResultSet): """Cursor for handling Polars DataFrame results from Athena queries. This cursor returns query results as Polars DataFrames using Polars' native diff --git a/pyathena/result_set.py b/pyathena/result_set.py index bfca7c58..b6c267bb 100644 --- a/pyathena/result_set.py +++ b/pyathena/result_set.py @@ -794,17 +794,22 @@ def _get_rows( ] -class WithResultSet: - """Mixin providing a SQL cursor's result set, its properties, and default fetch. +class WithResultSet(BaseCursor, CursorIterator): + """Base class of the SQL cursors that keep a result set. - Shared by the sync (``WithFetch``) and asyncio (``WithAsyncFetch``) cursors. - The fetch methods call the result set synchronously; the asyncio cursors - override them with async versions. + Provides the result set and its properties, fetch, ``close``, + ``executemany``, ``cancel``, and sync iteration. The sync SQL cursors + subclass it directly; ``WithAsyncFetch`` overrides the fetch and lifecycle + methods with async versions for the asyncio cursors. """ - def __init__(self): - """Initialize the mixin with no query ID and no result set.""" - super().__init__() + def __init__(self, **kwargs) -> None: + """Initialize the cursor with no query ID and no result set. + + Args: + **kwargs: Arguments passed to ``BaseCursor.__init__``. + """ + super().__init__(**kwargs) self._query_id: str | None = None self._result_set: AthenaResultSet | None = None @@ -1114,18 +1119,6 @@ def fetchall( result_set = cast(AthenaResultSet, self.result_set) return result_set.fetchall() - -class WithFetch(BaseCursor, WithResultSet, CursorIterator): - """Mixin providing sync lifecycle and iteration for SQL cursors. - - Adds ``close``, ``executemany``, and ``cancel`` to the properties and - default fetch methods of ``WithResultSet``, and the sync iteration - protocol through ``CursorIterator``. - - Subclasses override ``execute()`` and optionally ``__init__`` and - format-specific helpers. - """ - def close(self) -> None: """Close the cursor and release associated resources.""" self._rowcount = -1 diff --git a/pyathena/s3fs/cursor.py b/pyathena/s3fs/cursor.py index 6052b66e..27299ff5 100644 --- a/pyathena/s3fs/cursor.py +++ b/pyathena/s3fs/cursor.py @@ -8,14 +8,14 @@ from pyathena.error import OperationalError from pyathena.model import AthenaQueryExecution from pyathena.options import ExecuteOptions -from pyathena.result_set import WithFetch +from pyathena.result_set import WithResultSet from pyathena.s3fs.converter import DefaultS3FSTypeConverter from pyathena.s3fs.result_set import AthenaS3FSResultSet, CSVReaderType _logger = logging.getLogger(__name__) -class S3FSCursor(WithFetch): +class S3FSCursor(WithResultSet): """Cursor for reading CSV results via S3FileSystem without pandas/pyarrow. This cursor uses Python's standard csv module and PyAthena's S3FileSystem From 1b594b13f2977cb5b165b5c767ac11b63eecaf99 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Tue, 29 Sep 2026 23:52:41 +0900 Subject: [PATCH 6/8] Describe what WithAsyncFetch overrides in the WithResultSet docstring Co-Authored-By: Claude Opus 5.5 --- pyathena/result_set.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/pyathena/result_set.py b/pyathena/result_set.py index b6c267bb..eaf6f6b7 100644 --- a/pyathena/result_set.py +++ b/pyathena/result_set.py @@ -799,8 +799,9 @@ class WithResultSet(BaseCursor, CursorIterator): Provides the result set and its properties, fetch, ``close``, ``executemany``, ``cancel``, and sync iteration. The sync SQL cursors - subclass it directly; ``WithAsyncFetch`` overrides the fetch and lifecycle - methods with async versions for the asyncio cursors. + subclass it directly. For the asyncio cursors, ``WithAsyncFetch`` + overrides ``executemany`` and ``cancel`` with async versions, and its + subclasses override the fetch methods. """ def __init__(self, **kwargs) -> None: From a874aedaf00db23c7ee63799808c1cfd1d50a7f9 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Wed, 30 Sep 2026 00:04:14 +0900 Subject: [PATCH 7/8] Document the result set properties of WithResultSet Co-Authored-By: Claude Opus 5.5 --- pyathena/result_set.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/pyathena/result_set.py b/pyathena/result_set.py index eaf6f6b7..25b82ab6 100644 --- a/pyathena/result_set.py +++ b/pyathena/result_set.py @@ -823,6 +823,11 @@ def _reset_state(self) -> None: @property def result_set(self) -> AthenaResultSet | None: + """The result set of the last executed query. + + Returns: + The result set, or None before a query succeeds or after a reset. + """ return self._result_set @result_set.setter @@ -855,6 +860,12 @@ def catalog(self) -> str | None: @property def query_id(self) -> str | None: + """The query execution ID of the last started query. + + Returns: + The query execution ID, or None if no query has started since the + last reset. + """ return self._query_id @query_id.setter @@ -1057,6 +1068,13 @@ def rowcount(self) -> int: @property def arraysize(self) -> int: + """The number of rows that ``fetchmany()`` fetches by default. + + Setting it to zero or a negative value raises ``ProgrammingError``. + + Returns: + The default number of rows per ``fetchmany()`` call. + """ return self._arraysize @arraysize.setter @@ -1067,6 +1085,12 @@ def arraysize(self, value: int) -> None: @property def rownumber(self) -> int | None: + """The zero-based index of the next row in the result set. + + Returns: + The row index, or None if there is no result set or the index is + unknown. + """ return self.result_set.rownumber if self.result_set else None def fetchone( From 3e376d2d63c107391cb067c8f56dfb527934d267 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Wed, 30 Sep 2026 00:06:33 +0900 Subject: [PATCH 8/8] Correct the query_id and arraysize property docstrings Co-Authored-By: Claude Opus 5.5 --- pyathena/result_set.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/pyathena/result_set.py b/pyathena/result_set.py index 25b82ab6..fffbb43a 100644 --- a/pyathena/result_set.py +++ b/pyathena/result_set.py @@ -860,11 +860,14 @@ def catalog(self) -> str | None: @property def query_id(self) -> str | None: - """The query execution ID of the last started query. + """The query execution ID of the last execution. + + With ``cache_size`` or ``cache_expiration_time``, this can be the ID of + a previous execution whose result is reused. Returns: - The query execution ID, or None if no query has started since the - last reset. + The query execution ID, or None if there is none since the last + reset. """ return self._query_id @@ -1068,9 +1071,11 @@ def rowcount(self) -> int: @property def arraysize(self) -> int: - """The number of rows that ``fetchmany()`` fetches by default. + """The default number of rows per ``fetchmany()`` call. - Setting it to zero or a negative value raises ``ProgrammingError``. + ``execute()`` passes it to the new result set, so a change applies to + the result sets of later executions. Setting it to zero or a negative + value raises ``ProgrammingError``. Returns: The default number of rows per ``fetchmany()`` call.