Skip to content

fix(mlops): stop waiting on CANCELLED Athena queries in wait_for_athena_query - #6417

Open
AkshayShah03 wants to merge 1 commit into
aws:masterfrom
AkshayShah03:fix/athena-wait-cancelled
Open

AkshayShah03 wants to merge 1 commit into
aws:masterfrom
AkshayShah03:fix/athena-wait-cancelled

Conversation

@AkshayShah03

Copy link
Copy Markdown

Issue #, if available: N/A

Description of changes:

wait_for_athena_query only stops polling on SUCCEEDED or FAILED. CANCELLED is also a terminal Athena query state. Waiting on a query cancelled from the console, by a workgroup limit, or by StopQueryExecution polls every 5 seconds forever, which hangs AthenaQuery.wait(), run_athena_query() and the DatasetBuilder flows. v2 exited on any state other than QUEUED or RUNNING.

This change treats CANCELLED as terminal. run_athena_query and AthenaQuery already raise for any state other than SUCCEEDED, so a cancelled query now surfaces as an error.

Testing:

  • New TestWaitForAthenaQuery.test_returns_on_terminal_state, parametrized over SUCCEEDED, FAILED and CANCELLED. The CANCELLED case fails before this change.
  • The feature store feature_utils, athena_query and dataset_builder tests: 120 passed.

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

…na_query

wait_for_athena_query only stopped polling on SUCCEEDED or FAILED. CANCELLED is
also a terminal Athena state, so waiting on a cancelled query (from the
console, a workgroup limit, or StopQueryExecution) looped forever, hanging
AthenaQuery.wait(), run_athena_query() and DatasetBuilder flows.

Treat CANCELLED as terminal. Callers already raise for any state other than
SUCCEEDED.

This branch had an error being deployed

1 failed deployment
manual-approval — 5c239256 Deployed Oct 10, 2026 by AkshayShah03 via wait-for-approval #1969
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant