amoghrajesh opened a new pull request, #65590:
URL: https://github.com/apache/airflow/pull/65590
<!-- SPDX-License-Identifier: Apache-2.0
https://www.apache.org/licenses/LICENSE-2.0 -->
<!--
Thank you for contributing!
Please provide above a brief description of the changes made in this pull
request.
Write a good git commit message following this guide:
http://chris.beams.io/posts/git-commit/
Please make sure that your code changes are covered with tests.
And in case of new features or big changes remember to adjust the
documentation.
Feel free to ping (in general) for the review if you do not see reaction for
a few days
(72 Hours is the minimum reaction time you can expect from volunteers) - we
sometimes miss notifications.
In case of an existing issue, reference it using one of the following:
* closes: #ISSUE
* related: #ISSUE
-->
---
##### Was generative AI tooling used to co-author this PR?
<!--
If generative AI tooling has been used in the process of authoring this PR,
please
change below checkbox to `[X]` followed by the name of the tool, uncomment
the "Generated-by".
-->
- [x] Yes: Claude code
<!--
Generated-by: [Tool Name] following [the
guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
-->
closes: https://github.com/apache/airflow/issues/36090
### Why + What
Deferred tasks interacting with external systems (BigQuery, Databricks,
etc.) previously had no standard way to cancel the external job when a user
marked the task failed, cleared it, or marked it succeeded. This left orphaned
jobs running in the background.
To tackle this, adding three methods to `BaseTrigger`:
- `on_cancel()` — plug point to cancel the external job; no-op by default
- `safe_to_cancel()` — returns `True` only when the task state is not
DEFERRED, distinguishing user-initiated cancellation from triggerer
restart/redistribution by using task execution API
- `get_task_state()` — queries and gets the current task state via the
execution API
The `triggerer_job_runner` now calls `safe_to_cancel()` + `on_cancel()`
inside the existing `asyncio.CancelledError` handler, with
`suppress(BaseException)` so that `on_cancel()` failures do not prevent the
outer `CancelledError` from propagating.
### Impact
1. Existing trigger subclasses
Anyone who already implemented `get_task_state()` or `safe_to_cancel()` on
their own trigger subclass (like BigQuery) will need to rework it.
2. During normal triggerer shutdown or redistribution, every active trigger
fires CancelledError. The framework now calls safe_to_cancel() for each, which
makes a synchronous Execution API call via `sync_to_async`. Under high trigger
counts this can add latency to shutdown.
3. Provider authors / external triggers
Purely additive. `on_cancel()` defaults to a no-op so nothing breaks if you
do not implement it. Anyone who wants to hook into user-initiated cancellation
can override it.
### Testing
Trigger
```python
log = logging.getLogger(__name__)
SENTINEL_FILE = "/tmp/on_cancel_was_called"
class VerifyOnCancelTrigger(BaseTrigger):
"""
Trigger that sleeps forever.
on_cancel() writes a sentinel file so we can verify it was called
without needing a real external service like Databricks or BigQuery.
"""
def serialize(self) -> tuple[str, dict]:
return ("verify_on_cancel_trigger.VerifyOnCancelTrigger", {})
async def run(self) -> AsyncIterator[TriggerEvent]:
log.info("VerifyOnCancelTrigger running — mark this task as failed
to test on_cancel()")
await asyncio.sleep(99999)
yield TriggerEvent("done") # never reached
async def on_cancel(self) -> None:
log.info("VerifyOnCancelTrigger.on_cancel() CALLED — simulated
external job cancelled")
with open(SENTINEL_FILE, "w") as f:
f.write("cancelled")
```
DAG:
```python
class DeferForeverOperator(BaseOperator):
def execute(self, context):
self.defer(trigger=VerifyOnCancelTrigger(), method_name="resume")
def resume(self, context, event=None):
log.info("Resumed after trigger fired (should not happen in cancel
test)")
with DAG(
dag_id="test_on_cancel",
start_date=datetime(2024, 1, 1),
schedule=None,
tags=["manual-test"],
) as dag:
DeferForeverOperator(task_id="deferred_task")
```
The trigger is such that if `on_cancel` is called, once I mark it failed
from UI, it will create a sentinel file.
1. Cancelled trigger
<img width="1666" height="899" alt="image"
src="https://github.com/user-attachments/assets/ce095446-b002-443f-b811-77b0809d4280";
/>
2. Before run
```shell script
root@0c61b077a443:/opt/airflow$ ls /tmp/on_cancel_was_called
ls: cannot access '/tmp/on_cancel_was_called': No such file or directory
[Breeze:3.10.19] root@0c61b077a443:/opt/airflow$
[Breeze:3.10.19] root@0c61b077a443:/opt/airflow$
```
3. After run
```shell script
root@0c61b077a443:/opt/airflow$
[Breeze:3.10.19] root@0c61b077a443:/opt/airflow$ ls /tmp/on_cancel_was_called
/tmp/on_cancel_was_called
[Breeze:3.10.19] root@0c61b077a443:/opt/airflow$
[Breeze:3.10.19] root@0c61b077a443:/opt/airflow$
[Breeze:3.10.19] root@0c61b077a443:/opt/airflow$ cat
/tmp/on_cancel_was_called
cancelled[Breeze:3.10.19] root@0c61b077a443:/opt/airflow
```
---
* Read the **[Pull Request
Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)**
for more information. Note: commit author/co-author name and email in commits
become permanently public when merged.
* For fundamental code changes, an Airflow Improvement Proposal
([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals))
is needed.
* When adding dependency, check compliance with the [ASF 3rd Party License
Policy](https://www.apache.org/legal/resolved.html#category-x).
* For significant user-facing changes create newsfragment:
`{pr_number}.significant.rst`, in
[airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments).
You can add this file in a follow-up commit after the PR is created so you
know the PR number.
--
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]
