Recover Athena queries after worker crashes by 1fanwang · Pull Request #72419 · apache/airflow

1fanwang · 2026-09-02T06:10:13Z

A synchronous Athena task submits another query when its worker exits after StartQueryExecution and Airflow retries it. The retry now resumes the stored query. A completed query returns without resubmission; failed, cancelled, or missing queries submit normally.

When Athena no longer recognizes the stored ID, it returns InvalidRequestException: QueryExecution <id> was not found. Only that response for the same query ID triggers replacement; other invalid requests still fail.

AIP-103 task state is used only for synchronous polling. The operator stores the query ID before polling, reconnects active queries, and restores the ID and results link. Deferrable tasks and client_request_token behavior do not change.

Testing

uv run --project providers/amazon pytest providers/amazon/tests/unit/amazon/aws/operators/test_athena.py -xq

The live test drove the real operator execute() path and task state API against fresh SQLite with a local Athena hook stand-in.

UPSTREAM=/tmp/airflow-athena-main
git worktree add --detach "$UPSTREAM" upstream/main
AIRFLOW_HOME="$PWD/.tmp-athena-e2e" \
LEGACY_ATHENA_OPERATOR_PATH="$UPSTREAM/providers/amazon/src/airflow/providers/amazon/aws/operators/athena.py" \
uv run --project providers/amazon pytest -p tests_common.pytest_plugin \
  dev/test_athena_resumability_e2e.py -xvs
git worktree remove "$UPSTREAM"

Raw logs

TASK_STATE_AFTER_CRASH api=query-1 db=query-1
BEFORE stored=query-1 result=query-2 submitted=['query-2']
Reconnecting to existing job external_id=query-1 external_id_key=athena_query_execution_id status=RUNNING
AFTER stored=query-1 result=query-1 submitted=[] operator_id=query-1 link_id=query-1
PASSED
1 passed, 1 warning in 50.94s

Was generative AI tooling used to co-author this PR?

Yes (GitHub Copilot CLI (GPT-5.6 Sol))

Generated-by: GitHub Copilot CLI (GPT-5.6 Sol) following the guidelines

Synchronous Athena tasks currently submit a second query when a worker crashes after submission but before polling completes. This can duplicate work and leave the original query running. Generated-by: GitHub Copilot CLI (GPT-5.6 Sol) Signed-off-by: 1fanwang <[email protected]>

Stored query IDs can disappear before Airflow retries the task. Athena reports that lookup as InvalidRequestException, which otherwise prevents the retry from submitting a replacement. Generated-by: GitHub Copilot CLI (GPT-5.6 Sol) Signed-off-by: 1fanwang <[email protected]>

potiuk · 2026-09-25T20:22:42Z

Hello @1fanwang - thank you for your contributions to Apache Airflow!

The Airflow community has introduced a limit of 5 open pull requests at a time for contributors without write access to the repository. You currently have 34 open pull requests, so - as a one-time step of introducing the limit - we closed the ones where maintainers have not engaged yet:

Report failure causes and bound infrastructure replacements #66405 - Report failure causes and bound infrastructure replacements
Honor AirflowFailException raised inside on_retry_callback #66781 - Honor AirflowFailException raised inside on_retry_callback
Emit OpenTelemetry spans from airflow CLI entry points #66789 - Emit OpenTelemetry spans from airflow CLI entry points
Preserve sensor start_date when scheduler advances reschedule to queued #66790 - Preserve sensor start_date when scheduler advances reschedule to queued
Reuse the resumable core for TriggerDagRunOperator's durable wait #68955 - Reuse the resumable core for TriggerDagRunOperator's durable wait
Let TriggerDagRunOperator own its execution via a new accessor #69135 - Let TriggerDagRunOperator own its execution via a new accessor
Add resumable_job.reconnect_failure metric for failed reconnects #70653 - Add resumable_job.reconnect_failure metric for failed reconnects
Allow configuring the collation of asset name columns #71249 - Allow configuring the collation of asset name columns
Fix task_instance_history downgrade failing under the PyMySQL driver #71258 - Fix task_instance_history downgrade failing under the PyMySQL driver
Correct the retry_db_transaction docstring about provide_session #71311 - Correct the retry_db_transaction docstring about provide_session
Show worker pod logs for tasks that failed before writing any log #71566 - Show worker pod logs for tasks that failed before writing any log
Reconnect Synapse Spark batches on task retry #72368 - Reconnect Synapse Spark batches on task retry
Reconnect Azure Data Factory pipelines on task retry #72399 - Reconnect Azure Data Factory pipelines on task retry
Prevent duplicate Anthropic batches on task retry #72400 - Prevent duplicate Anthropic batches on task retry
Reconnect AWS Batch jobs after worker crashes #72413 - Reconnect AWS Batch jobs after worker crashes
Reconnect Airbyte sync jobs after task retries #72414 - Reconnect Airbyte sync jobs after task retries
Reconnect Data Fusion pipelines after task retries #72415 - Reconnect Data Fusion pipelines after task retries
Prevent duplicate OpenAI batches on task retry #72416 - Prevent duplicate OpenAI batches on task retry
Recover Dataproc jobs after worker crashes #72417 - Recover Dataproc jobs after worker crashes
Recover Athena queries after worker crashes #72419 - Recover Athena queries after worker crashes
Reconnect Cloud Batch jobs on task retry #72420 - Reconnect Cloud Batch jobs on task retry
Prevent duplicate ECS submissions on task retry #72421 - Prevent duplicate ECS submissions on task retry
Recover Databricks SQL statements after task retries #72422 - Recover Databricks SQL statements after task retries
Reconnect dbt Cloud runs after worker crashes #72423 - Reconnect dbt Cloud runs after worker crashes
Restore missing serialized Dags without a full reparse #72927 - Restore missing serialized Dags without a full reparse
Report task failure causes without changing retries #73350 - Report task failure causes without changing retries

These pull requests stay open because maintainers are already engaged in them - they count towards your limit:

Add durable option to TriggerDagRunOperator to reconnect on retry #68936 - Add durable option to TriggerDagRunOperator to reconnect on retry
Let deferrable tasks reconnect without polling #68952 - Let deferrable tasks reconnect without polling
Keep database failures from looking like missing secrets #71307 - Retry transient database errors in the metastore secrets backend
Return each attempt's own URL from operator extra links #71522 - Return each attempt's own URL from operator extra links
Return 404 when a DagRun disappears during task API writes #72907 - Return 404 instead of 500 when DagRun is missing in Execution API
Edge: Return 404 for missing task logs #72908 - Edge: Return 404 for missing task logs and emit failed job metrics
Edge: Return 400 instead of 500 on invalid worker version #72925 - Edge: Return 400 instead of 500 on invalid worker version
Return 404 for task-state writes after a Dag run is deleted #72926 - Return 404 for task-state writes after a Dag run is deleted

This is not a judgement of you or of your changes. We never told contributors before that opening many pull requests at once was a problem, so there is nothing to feel bad about - and nothing is lost: your branches, commits and the review history stay where they are.

What we ask you to do is to make your first prioritization decision: choose which of the pull requests above matter most to you, and reopen them (up to 5 open at a time, including the ones still open) with the "Reopen pull request" button or gh pr reopen <PR_NUMBER> --repo apache/airflow. Reopen the ones you are ready to follow through - keep them rebased, respond to review comments and fix failing checks.

While your pull requests are waiting for review, the most valuable thing you can do is help in other ways - reviewing other contributors' pull requests, helping with issues, and taking part in the discussions on the devlist and Slack.

Why we introduced the limit, what it means for you and how to reopen or restore a pull request is explained in https://github.com/apache/airflow/blob/main/contributing-docs/32_open_pull_request_limit.rst.

Drafted-by: Claude Code (Opus 5); reviewed by @potiuk before posting

boring-cyborg Bot added area:providers provider:amazon AWS/Amazon - related issues labels Sep 2, 2026

potiuk added the closed because of open PR limit Closed as a one-time step of introducing the open pull request limit label Sep 25, 2026

potiuk closed this Sep 25, 2026

This was referenced Sep 25, 2026

Reconnect Cloud Batch jobs on task retry #72420

Closed

Prevent duplicate ECS submissions on task retry #72421

Closed

Recover Databricks SQL statements after task retries #72422

Closed

Reconnect dbt Cloud runs after worker crashes #72423

Closed

This was referenced Sep 25, 2026

Restore missing serialized Dags without a full reparse #72927

Closed

Report task failure causes without changing retries #73350

Closed

Navigation Menu

Repository navigation

Recover Athena queries after worker crashes - #72419

Recover Athena queries after worker crashes#72419
1fanwang wants to merge 2 commits into
apache:mainfrom
1fanwang:recover-athena-queries

1fanwang commented Sep 2, 2026 •

edited

Loading

Uh oh!

potiuk commented Sep 25, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

1fanwang commented Sep 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Testing

Was generative AI tooling used to co-author this PR?

Uh oh!

potiuk commented Sep 25, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

1fanwang commented Sep 2, 2026 •

edited

Loading