Skip to content

Commit d79d2ae

Browse files
Add docstrings and fix log message in aio cursor
- Add Google-style docstrings to AioCursor properties and methods matching the sync Cursor docstring patterns - Add docstrings to AthenaAioResultSet fetch methods - Fix _poll log message to match sync: "Query canceled by user." Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1 parent 7bdc3c1 commit d79d2ae

File tree

3 files changed

+104
-1
lines changed

3 files changed

+104
-1
lines changed

pyathena/aio/common.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -103,7 +103,7 @@ async def _poll(self, query_id: str) -> AthenaQueryExecution: # type: ignore[ov
103103
query_execution = await self.__poll(query_id)
104104
except asyncio.CancelledError:
105105
if self._kill_on_interrupt:
106-
_logger.warning("Query canceled.")
106+
_logger.warning("Query canceled by user.")
107107
await self._cancel(query_id)
108108
query_execution = await self.__poll(query_id)
109109
else:

pyathena/aio/cursor.py

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,12 @@ def __init__(
6363

6464
@property
6565
def result_set(self) -> Optional[AthenaResultSet]:
66+
"""Get the result set from the last executed query.
67+
68+
Returns:
69+
The result set object containing query results, or None if no
70+
query has been executed or the query didn't return results.
71+
"""
6672
return self._result_set
6773

6874
@result_set.setter
@@ -71,6 +77,12 @@ def result_set(self, val) -> None:
7177

7278
@property
7379
def query_id(self) -> Optional[str]:
80+
"""Get the Athena query execution ID of the last executed query.
81+
82+
Returns:
83+
The query execution ID assigned by Athena, or None if no query
84+
has been executed.
85+
"""
7486
return self._query_id
7587

7688
@query_id.setter
@@ -79,13 +91,32 @@ def query_id(self, val) -> None:
7991

8092
@property
8193
def rownumber(self) -> Optional[int]:
94+
"""Get the current row number within the result set.
95+
96+
Returns:
97+
The zero-based index of the current row, or None if no result set
98+
is available or no rows have been fetched.
99+
"""
82100
return self.result_set.rownumber if self.result_set else None
83101

84102
@property
85103
def rowcount(self) -> int:
104+
"""Get the number of rows affected by the last operation.
105+
106+
For SELECT statements, this returns the total number of rows in the
107+
result set. For other operations, behavior follows DB API 2.0 specification.
108+
109+
Returns:
110+
The number of rows, or -1 if not applicable or unknown.
111+
"""
86112
return self.result_set.rowcount if self.result_set else -1
87113

88114
def close(self) -> None:
115+
"""Close the cursor and free any associated resources.
116+
117+
Closes the cursor and any associated result sets. This method is
118+
synchronous (no I/O needed, just clears references).
119+
"""
89120
if self.result_set and not self.result_set.is_closed:
90121
self.result_set.close()
91122

@@ -158,18 +189,47 @@ async def executemany( # type: ignore[override]
158189
seq_of_parameters: List[Optional[Union[Dict[str, Any], List[str]]]],
159190
**kwargs,
160191
) -> None:
192+
"""Execute a SQL query multiple times with different parameters.
193+
194+
Args:
195+
operation: SQL query string to execute.
196+
seq_of_parameters: Sequence of parameter dictionaries or lists,
197+
one for each execution.
198+
**kwargs: Additional keyword arguments passed to each
199+
``execute()`` call.
200+
201+
Note:
202+
Operations that return result sets are not allowed with
203+
``executemany``.
204+
"""
161205
for parameters in seq_of_parameters:
162206
await self.execute(operation, parameters, **kwargs)
207+
# Operations that have result sets are not allowed with executemany.
163208
self._reset_state()
164209

165210
async def cancel(self) -> None:
211+
"""Cancel the currently executing query.
212+
213+
Raises:
214+
ProgrammingError: If no query is currently executing
215+
(query_id is None).
216+
"""
166217
if not self.query_id:
167218
raise ProgrammingError("QueryExecutionId is none or empty.")
168219
await self._cancel(self.query_id)
169220

170221
async def fetchone( # type: ignore[override]
171222
self,
172223
) -> Optional[Union[Any, Dict[Any, Optional[Any]]]]:
224+
"""Fetch the next row of a query result set.
225+
226+
Returns:
227+
A tuple representing the next row, or None if no more rows.
228+
229+
Raises:
230+
ProgrammingError: If called before executing a query that
231+
returns results.
232+
"""
173233
if not self.has_result_set:
174234
raise ProgrammingError("No result set.")
175235
result_set = cast(AthenaAioResultSet, self.result_set)
@@ -178,6 +238,18 @@ async def fetchone( # type: ignore[override]
178238
async def fetchmany( # type: ignore[override]
179239
self, size: Optional[int] = None
180240
) -> List[Union[Any, Dict[Any, Optional[Any]]]]:
241+
"""Fetch multiple rows from a query result set.
242+
243+
Args:
244+
size: Maximum number of rows to fetch. If None, uses arraysize.
245+
246+
Returns:
247+
List of tuples representing the fetched rows.
248+
249+
Raises:
250+
ProgrammingError: If called before executing a query that
251+
returns results.
252+
"""
181253
if not self.has_result_set:
182254
raise ProgrammingError("No result set.")
183255
result_set = cast(AthenaAioResultSet, self.result_set)
@@ -186,6 +258,15 @@ async def fetchmany( # type: ignore[override]
186258
async def fetchall( # type: ignore[override]
187259
self,
188260
) -> List[Union[Any, Dict[Any, Optional[Any]]]]:
261+
"""Fetch all remaining rows from a query result set.
262+
263+
Returns:
264+
List of tuples representing all remaining rows in the result set.
265+
266+
Raises:
267+
ProgrammingError: If called before executing a query that
268+
returns results.
269+
"""
189270
if not self.has_result_set:
190271
raise ProgrammingError("No result set.")
191272
result_set = cast(AthenaAioResultSet, self.result_set)

pyathena/aio/result_set.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -147,6 +147,14 @@ async def _async_pre_fetch(self) -> None:
147147
async def fetchone( # type: ignore[override]
148148
self,
149149
) -> Optional[Union[Tuple[Optional[Any], ...], Dict[Any, Optional[Any]]]]:
150+
"""Fetch the next row of the result set.
151+
152+
Automatically fetches the next page from Athena when the current
153+
page is exhausted and more pages are available.
154+
155+
Returns:
156+
A tuple representing the next row, or None if no more rows.
157+
"""
150158
if not self._rows and self._next_token:
151159
await self._async_fetch()
152160
if not self._rows:
@@ -159,6 +167,15 @@ async def fetchone( # type: ignore[override]
159167
async def fetchmany( # type: ignore[override]
160168
self, size: Optional[int] = None
161169
) -> List[Union[Tuple[Optional[Any], ...], Dict[Any, Optional[Any]]]]:
170+
"""Fetch multiple rows from the result set.
171+
172+
Args:
173+
size: Maximum number of rows to fetch. If None, uses arraysize.
174+
175+
Returns:
176+
List of row tuples. May contain fewer rows than requested if
177+
fewer are available.
178+
"""
162179
if not size or size <= 0:
163180
size = self._arraysize
164181
rows = []
@@ -173,6 +190,11 @@ async def fetchmany( # type: ignore[override]
173190
async def fetchall( # type: ignore[override]
174191
self,
175192
) -> List[Union[Tuple[Optional[Any], ...], Dict[Any, Optional[Any]]]]:
193+
"""Fetch all remaining rows from the result set.
194+
195+
Returns:
196+
List of all remaining row tuples.
197+
"""
176198
rows = []
177199
while True:
178200
row = await self.fetchone()

0 commit comments

Comments
 (0)