rebenitez1802 opened a new pull request, #45860:
URL: https://github.com/apache/superset/pull/45860
### SUMMARY
When an async chart-data task is aborted or times out, it asks the database
to cancel its query. The cancel ID was captured once, when the task's first
cursor opened, so for Impala it cancelled the wrong operation, or nothing:
- impyla runs `USE <db>` while creating a cursor for a URL that names a
database. The ID read at that point belongs to that finished `USE`. The abort
cancelled it (the coordinator answers HTTP 200), and the chart's SELECT ran to
completion.
- With no database in the URL there was no ID at all, so the task could not
be aborted.
- On every engine, only the first ID was kept, so later statements and later
cursors (a time comparison's offset query, for example) were never targeted.
**Changes**
- **Phases.** `get_df` notifies the task's cursor sink with the cursor, its
database and a phase: opened, submitting, submitted, then released or failed.
`capture_cancel_query_id`, `cancel_chart_query` and MCP's cancellation hooks
are unchanged.
- **The task follows the running operation.** Where the ID is read depends
on `has_query_id_before_execute`:
- **Session ID** (PostgreSQL, MySQL, …): read once per cursor, when it
opens. It is never read after execute, because these engines read it by running
SQL on the live cursor, which would replace the chart's pending rows.
- **Operation ID** (Impala): read after each statement is submitted, and
retired before the next one.
- **Only the task's own database.** A cursor on another database, such as an
annotation layer's source chart, never replaces the task's handle.
- **Abortable from cursor open.** Impala registers the abort handler when
the cursor opens. An abort that arrives while a statement is being submitted is
recorded, and the worker cancels the operation once its ID exists, because
abort handlers run only once.
- **No further work after an abort.** An aborting task starts no further
statement and fetches no further results. It raises `ChartQueryAbortedError`, a
`SupersetErrorException`, so the datasource's `query()` re-raises it instead of
returning a failed result. This matters for a time comparison: an aborted
offset query must not be merged as empty and the chart cached without its
comparison series. A cursor whose statements fail while the task is aborting
reports the abort, not the driver's error.
- **The orphan reaper's handle follows the operation.** It is replaced as
the operation changes and cleared when a cursor is released, through
`TaskContext.persist_cancellation`. A timeout's "Task timed out" message is
kept through these writes.
- **The abort handler uses its own session.** It loads the database in its
own thread's session instead of using the worker's instance.
The transport is unchanged. Impala still cancels through the coordinator web
UI, which needs `IMPALA_CANCEL_QUERY_ALLOW_INTERNAL_HOSTS` for internal hosts.
This is noted in `UPDATING.md`.
Open PRs #44406 and #44552 change the same cursor seam. #44406's database
filter is covered here. If #44552 lands first, its Explore Stop sink needs to
accept the phase argument and publish at `submitted`.
### BEFORE/AFTER SCREENSHOTS OR ANIMATED GIF
Not applicable; there are no UI changes. The coordinator results for an
abort on Impala 4.5.2 (`/queries` and the query profile):
| Case | Before | After |
| --- | --- | --- |
| URL with a database, abort while running | cancel sent to the finished
`USE`; SELECT ran to completion (10.5 s) | SELECT cancelled at about 2.4 s |
| URL without a database, abort while running | task not abortable; SELECT
ran to completion | SELECT cancelled at about 2.3 s |
| Abort while the statement is being submitted | `USE` targeted, or not
abortable | the worker cancels the SELECT as soon as it has an ID (cancelled
within 0.15 s of starting) |
| Time comparison, abort during the offset query | the first cursor's `USE`
targeted; offset query ran to completion | offset query cancelled at about 2.1
s; the task raises instead of caching a partial result |
### TESTING INSTRUCTIONS
**Unit tests.** They use fake impyla-like cursors. The PostgreSQL ordering
guard and the time-comparison cache test use real SQLite cursors and a real
`SqlaTable`.
```
pytest tests/unit_tests/tasks/test_query_cancel.py
tests/unit_tests/tasks/test_context.py \
tests/unit_tests/tasks/test_reap.py tests/unit_tests/daos/test_tasks.py
```
**Real Impala.** These runs used the public Apache Impala 4.5.2 quickstart.
They drive the task's cancellation and `get_df` through the real task abort
path, and read the results from the coordinator. They did not run a full
Explore request on a Celery worker.
### ADDITIONAL INFORMATION
- [ ] Has associated issue:
- [ ] Required feature flags:
- [ ] Changes UI
- [ ] Includes DB Migration (follow approval process in
[SIP-59](https://github.com/apache/superset/issues/13351))
- [ ] Migration is atomic, supports rollback & is backwards-compatible
- [ ] Confirm DB migration upgrade and downgrade tested
- [ ] Runtime estimates and downtime expectations provided
- [ ] Introduces new feature or API
- [ ] Removes existing feature or API
**Follow-up, not in this PR:** if the reaper reaps a worker that is still
alive, that worker's later handle write can overwrite the reaper's error
message. A write that merges only the task's private state under a row lock
would close that gap.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
