Fail sensor tasks immediately if the ExecutionMode.WATCHER producer task fails

[pankajastro]

relates to: #1960

This PR adds a default on_failure_callback to the producer task. The callback records the producer task’s state and ensures that the sensor fails immediately if the producer task has failed, improving error visibility and reducing unnecessary wait time.

[Copilot]

Pull Request Overview

This PR enhances the watcher functionality by making sensor tasks fail immediately when the producer task fails, rather than waiting for a timeout. It adds a default on_failure_callback to the producer task that records its failure state via XCom, which the sensor then checks before attempting to poll for events.

Key changes:

• Added default failure callback to producer task that pushes "failed" state to XCom
• Modified sensor's poke method to check producer task state and fail immediately if producer has failed
• Implemented version-aware callback handling to support both single callbacks (Airflow < 2.6.0) and callback lists (Airflow >= 2.6.0)

Reviewed Changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.

File	Description
cosmos/operators/watcher.py	Added producer task failure callback, state checking in sensor, and version-aware callback merging logic
tests/operators/test_watcher.py	Updated test mock to account for additional XCom pull for producer state check

---

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

[Copilot]

Pull Request Overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

---

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 97.71%. Comparing base (fa81082) to head (4d712a2).

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2040   +/-   ##
=======================================
  Coverage   97.71%   97.71%           
=======================================
  Files          89       89           
  Lines        5681     5700   +19     
=======================================
+ Hits         5551     5570   +19     
  Misses        130      130           

☔ 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.

[pankajastro]

I'll create a follow-up PR to extend the documentation introduced in PR #2046

[tatiana]

Thanks a lot for working on this, @pankajastro !

I have a behavioural question. In the case of Jaffle Shop, let's say there was an issue applying a seed. The producer task fails. Will all watcher tasks be marked as failed, even if they had not been scheduled and did not start sensing?

Please, could you increase the coverage?

[pankajastro]

I have a behavioural question. In the case of Jaffle Shop, let's say there was an issue applying a seed. The producer task fails. Will all watcher tasks be marked as failed, even if they had not been scheduled and did not start sensing?

If not scheduled, it should go to the upstream_failed state

[tatiana]

If not scheduled, it should go to the upstream_failed state

@pankajastro, Did you run a few DAGs to confirm this behaviour and ensure that tasks skipped because upstream tasks failed would not enter sensor mode when they attempt to run for the "first time"?

[pankajastro]

@pankajastro, Did you run a few DAGs to confirm this behaviour and ensure that tasks skipped because upstream tasks failed would not enter sensor mode when they attempt to run for the "first time"?

I'm not sure if I understand it fully, but why should it be skipped?

[tatiana]

@pankajastro In the case of Jaffle Shop, let's say the following happens:

The dbt_producer_watcher fails while trying to apply the seed raw_customers.
The task raw_customers, which was actively watching, will fail (it becomes red). It will have logs for a first try that was unsuccessful. If it is retried, that would be its second try.
However, the task stg_customers will have never been started; it would never have been initialised as a sensor. It will likely be marked as upstream_failed (I'm asking you to confirm this). If this hypothesis is correct, it will likely never be attempted, and therefore, no logs will be generated for this nonexistent first attempt.

So, my concern is—and this is what I'm asking you to test and confirm—what happens when someone clears the status of the seed raw_customers and its downstream tasks? Let's say it succeeds; the task stg_customers will then proceed to its second attempt, as expected, and execute dbt run. However, I'm worried that the try number on the customers task would be 1 (afterall, it is the first time it is actually being executed), and it would enter sensor mode instead of running the dbt run, and it would wait "forever" (until timeout happens) for an update from the dbt producer that will never come.

[tatiana]

@pankajastro In the case of Jaffle Shop, let's say the following happens:

The dbt_producer_watcher fails while trying to apply the seed raw_customers. The task raw_customers, which was actively watching, will fail (it becomes red). It will have logs for a first try that was unsuccessful. If it is retried, that would be its second try. However, the task stg_customers will have never been started; it would never have been initialised as a sensor. It will likely be marked as upstream_failed (I'm asking you to confirm this). If this hypothesis is correct, it will likely never be attempted, and therefore, no logs will be generated for this nonexistent first attempt.

So, my concern is—and this is what I'm asking you to test and confirm—what happens when someone clears the status of the seed raw_customers and its downstream tasks? Let's say it succeeds; the task stg_customers will then proceed to its second attempt, as expected, and execute dbt run. However, I'm worried that the try number on the customers task would be 1 (afterall, it is the first time it is actually being executed), and it would enter sensor mode instead of running the dbt run, and it would wait "forever" (until timeout happens) for an update from the dbt producer that will never come.

I implemented the approach that was discussed with @pankajastro @pankajkoti in #2062

[tatiana]

@pankajastro @pankajkoti and I discussed a few scenarios, and we are extending the original implementation in this PR in a separate PR #2062. We still believe more work will be needed in this area. It seems that Airflow does not expose whether a task has a State.UPSTREAM_FAILED - which affects how we'll handle scenarios.