Skip to content

Commit f4b8f64

Browse files
Document on_poll polling callback
Add a "Query polling callback" section to docs/usage.md covering on_poll: connection- and cursor-level configuration, the synchronous callback contract, per-iteration invocation including the terminal state, async cursor usage, and the Spark calculation-status payload. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
1 parent 7af6317 commit f4b8f64

1 file changed

Lines changed: 80 additions & 0 deletions

File tree

docs/usage.md

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -427,6 +427,86 @@ The `on_start_query_execution` callback is supported by the following cursor typ
427427
Note: `AsyncCursor` and its variants do not support this callback as they already
428428
return the query ID immediately through their different execution model.
429429

430+
## Query polling callback
431+
432+
PyAthena provides an `on_poll` callback that is invoked once per poll iteration with the
433+
current query execution object, while PyAthena waits for the query to finish. This is useful
434+
for rendering live query progress (state, elapsed time, data scanned) in interactive
435+
environments such as Jupyter notebooks.
436+
437+
The callback is optional (`None` by default), so there is no impact on existing behaviour or
438+
performance when it is not used. It must be a **synchronous** function with the signature
439+
`Callable[[AthenaQueryExecution], None]` (for Spark calculations it receives an
440+
`AthenaCalculationExecutionStatus`). It is invoked on every poll, including the final
441+
iteration that observes the terminal state (`SUCCEEDED`, `FAILED`, or `CANCELLED`).
442+
443+
Unlike `on_start_query_execution`, `on_poll` is configured at the connection or cursor level
444+
only (there is no execute-level override).
445+
446+
### Connection-level callback
447+
448+
```python
449+
from pyathena import connect
450+
451+
def on_poll(query_execution):
452+
print(
453+
f"State: {query_execution.state}, "
454+
f"scanned: {query_execution.data_scanned_in_bytes} bytes"
455+
)
456+
457+
cursor = connect(
458+
s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
459+
region_name="us-west-2",
460+
on_poll=on_poll,
461+
).cursor()
462+
463+
cursor.execute("SELECT * FROM many_rows") # on_poll is invoked on each poll
464+
```
465+
466+
### Cursor-level callback
467+
468+
```python
469+
from pyathena import connect
470+
471+
def on_poll(query_execution):
472+
print(f"State: {query_execution.state}")
473+
474+
conn = connect(
475+
s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
476+
region_name="us-west-2",
477+
)
478+
cursor = conn.cursor(on_poll=on_poll)
479+
cursor.execute("SELECT * FROM many_rows")
480+
```
481+
482+
### Asynchronous cursors
483+
484+
`on_poll` also works with the asynchronous cursors. Because polling runs in a background
485+
thread (or event loop), the callback runs there too, so keep it lightweight and thread-safe:
486+
487+
```python
488+
from pyathena import connect
489+
from pyathena.pandas.async_cursor import AsyncPandasCursor
490+
491+
def on_poll(query_execution):
492+
print(f"State: {query_execution.state}")
493+
494+
conn = connect(
495+
s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
496+
region_name="us-west-2",
497+
)
498+
cursor = conn.cursor(AsyncPandasCursor, on_poll=on_poll)
499+
query_id, future = cursor.execute("SELECT * FROM many_rows")
500+
result = future.result()
501+
```
502+
503+
### Supported cursor types
504+
505+
`on_poll` is supported by **all** cursor types, since polling is shared by the base cursor:
506+
the synchronous cursors, the `Async*` cursors, the native-async `Aio*` cursors, and the Spark
507+
cursors. For Spark cursors the callback receives the per-poll
508+
`AthenaCalculationExecutionStatus` rather than an `AthenaQueryExecution`.
509+
430510
## Type hints for complex types
431511

432512
*New in version 3.30.0.*

0 commit comments

Comments
 (0)