Skip to content

fix(athena-driver): propagate query cancellation to Athena via StopQueryExecution#10861

Merged
igorlukanin merged 7 commits into
masterfrom
igor/core-477-ability-to-propagate-query-cancelation-to-athena
May 15, 2026
Merged

fix(athena-driver): propagate query cancellation to Athena via StopQueryExecution#10861
igorlukanin merged 7 commits into
masterfrom
igor/core-477-ability-to-propagate-query-cancelation-to-athena

Conversation

@igorlukanin
Copy link
Copy Markdown
Member

Summary

  • The Athena driver never called StopQueryExecution, so when Cube cancelled or timed out a query the underlying Athena execution kept running (and billing) on AWS. The driver now participates in Cube's existing cancellation infrastructure via cancelCombinator.
  • Each driver entry point (query, memory, stream) and each internal helper that issues an Athena query (loadPreAggregationIntoTable, unloadWithSql, unloadWithTable, queryColumnTypes) returns a promise whose .cancel calls StopQueryExecution for the in-flight QueryExecutionId. The stream() release callback (previously a no-op marked with a TODO) now does the same.
  • Cancellation is best-effort: a failure to stop the query is logged via console.warn and never bubbles back to the caller — the query is already being abandoned.

Test plan

  • Two integration tests added to packages/cubejs-athena-driver/test/AthenaDriver.test.ts (CI has Athena credentials): one for query, one for stream. Both kick off a deliberately slow Athena query, call .cancel() on the returned promise, and then poll GetQueryExecution directly to assert Athena ended up in CANCELLED/FAILED (never SUCCEEDED/RUNNING/QUEUED).
  • yarn tsc --noEmit clean in packages/cubejs-athena-driver.
  • yarn lint clean in packages/cubejs-athena-driver.
  • CI green.

@igorlukanin igorlukanin requested a review from a team as a code owner May 11, 2026 17:39
@github-actions github-actions Bot added driver:athena Issues related to the AWS Athena driver javascript Pull requests that update Javascript code data source driver labels May 11, 2026
@claude
Copy link
Copy Markdown
Contributor

claude Bot commented May 11, 2026

Claude finished @igorlukanin's task in 2m 49s —— View job


Code Review: Athena Query Cancellation

  • Read and analyze changed files
  • Review code quality and correctness
  • Review security considerations
  • Review performance implications
  • Review test coverage
  • Post final review

Overall: This is a well-executed PR. The cancellation pattern is applied consistently across all entry points, follows the established cancelCombinator convention used by other Cube drivers, and the implementation is clean.

Code Quality & Correctness

Strengths:

  • The cancelCombinator + saveCancelFn + waitForSuccessCancellable pattern is applied uniformly to all 7 methods that issue Athena queries (query, memory, stream, loadPreAggregationIntoTable, unloadWithSql, unloadWithTable, queryColumnTypes). No entry point is missed.
  • stopQuery correctly uses best-effort semantics (try/catch with console.warn). This is the right call — the query is already being abandoned, so bubbling a secondary failure would be counterproductive.
  • The async → sync signature change (returning the cancelCombinator promise directly) is correct and necessary for the .cancel property to be visible to callers.

Minor observations (non-blocking, left as inline comments):

  1. stream() release callback — The comment says "stop the query so it stops billing," but by the time stream() resolves, waitForSuccess has already completed — the Athena query is in SUCCEEDED state. The release call to stopQuery is always a no-op. The real cancellation path is through cancelCombinator during the wait phase. Suggested a more accurate comment.
  2. unloadWithSql nested cancellation gapqueryColumnTypes() is awaited directly, not via the outer saveCancelFn. If the outer promise is cancelled during this sub-call, the LIMIT 0 query won't be stopped. Negligible in practice since those are near-instant.
  3. Inherent race in cancelCombinator pattern — Tiny window between startQuery and saveCancelFn where cancel is a no-op. Not introduced by this PR; just flagged for awareness.

Security

No concerns. StopQueryExecution uses the same Athena credentials and the QueryExecutionId originates from Athena's own response — no user-supplied input reaches the stop call.

Performance

Net positive. Previously, cancelled/timed-out queries continued running (and billing) on Athena. Now they're stopped promptly. The only added overhead is one StopQueryExecution API call per cancellation — trivial.

Testing

The two integration tests are well-designed:

  • Realistic slow query (100M-row aggregation)
  • Captures QueryExecutionId via monkey-patch with proper cleanup in finally
  • Polls Athena's own execution state to verify the query actually stopped — not just testing the local code path but the end-to-end AWS behavior
  • Covers both query() and stream() paths
  • Timeout bumped from 2min → 3min to accommodate the additional test time

One scenario not tested: cancellation during loadPreAggregationIntoTable or unloadWithSql. These follow the exact same pattern, so the coverage from query/stream gives reasonable confidence.

Summary

Solid, focused change. The inline comments are all non-blocking suggestions. Looks good to merge once CI is green.

Comment thread packages/cubejs-athena-driver/src/AthenaDriver.ts Outdated
Comment thread packages/cubejs-athena-driver/src/AthenaDriver.ts Outdated
Comment thread packages/cubejs-athena-driver/src/AthenaDriver.ts Outdated
Comment thread packages/cubejs-athena-driver/test/AthenaDriver.test.ts Outdated
@codecov
Copy link
Copy Markdown

codecov Bot commented May 11, 2026

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 58.36%. Comparing base (d5da121) to head (07e36a5).
⚠️ Report is 11 commits behind head on master.

❗ There is a different number of reports uploaded between BASE (d5da121) and HEAD (07e36a5). Click for more details.

HEAD has 1 upload less than BASE
Flag BASE (d5da121) HEAD (07e36a5)
cubesql 1 0
Additional details and impacted files
@@             Coverage Diff             @@
##           master   #10861       +/-   ##
===========================================
- Coverage   78.86%   58.36%   -20.50%     
===========================================
  Files         470      216      -254     
  Lines       92289    16944    -75345     
  Branches     3435     3435               
===========================================
- Hits        72784     9890    -62894     
+ Misses      19003     6552    -12451     
  Partials      502      502               
Flag Coverage Δ
cube-backend 58.36% <ø> (+<0.01%) ⬆️
cubesql ?

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@igorlukanin igorlukanin merged commit 233e86f into master May 15, 2026
117 of 119 checks passed
@igorlukanin igorlukanin deleted the igor/core-477-ability-to-propagate-query-cancelation-to-athena branch May 15, 2026 19:15
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

data source driver driver:athena Issues related to the AWS Athena driver javascript Pull requests that update Javascript code

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant