Fix ExecutionMode.WATCHER behaviour with DbtTaskGroup

[tatiana]

While validating TaskGroup in a session with an end-user with Cosmos v1.11.0a8, we identified that it was not working as expected.

There were two main issues:

1. We were not setting the producer_task_id argument correctly in the sensor tasks
2. The producer task was always being picked up after the sensor tasks, even though it had a higher priority weight and was an "upstream" task of the dbt root nodes (with a trigger rule always, so they could start sensing).

The first issue was a bug, which is now fixed.

The second issue occurs because the sensor tasks have a trigger_rule of "always". This setting works well in the case of DbtDag, as the only upstream task is the producer task, and the producer task and root dbt project nodes can be scheduled at the same time, with the task weight of the producer giving its precedence. However, in the case of a DbtTaskGroup with one or multiple upstream tasks, the producer task must wait for the upstream task, whereas the dbt project root nodes ("top sensors") do not have to wait. This leads to sensors occupying the available lots before the producer is scheduled to arrive.

At some point, we suspected this was an issue with Airflow itself, but after a troubleshooting session with Ash, we realised it was not a bug in Airflow: apache/airflow#56723.

The side-effect of the implementation in this PR is that DbtTaskGroup and DbtDag are behaving differently:

• In DbtDag, the producer is upstream of the dbt root nodes
• In DbtTaskGroup, the producer is not an upstream root node

The reason we are keeping this behaviour in DbtDag is that it is not leading to any issues. It is topologically "correct", making it easier for users to see the dependency between the tasks & potentially view the task list topologically sorted.

At this moment, we were unable to find a way to offer this same behaviour with DbtTaskGroup. @ashb , with whom I troubleshooted the original "problem", pointed out Airflow could potentially have a different "state" that wouldn't always be, but something like "start when my upstream task started". However, by examining Airflow's current implementation, any changes to this area could potentially incur an additional cost to the scheduler and compromise Airflow's performance.

If we were able to "force" in the Airflow UI to have the producer task on the "top" both in the graph view and in the list of functions view, we could make DbtDagand DbtTaskGroup more consistent (by not having the producer task as upstream and dbt root nodes with trigger_rule always). Currently, in Airflow 3.0.x and Airflow 3.1.0, the producer task has been showing randomly, unfortunately, but we will work with the Airflow team to improve this and add follow-up tickets as needed.

Closes: #1961
Closes: https://github.com/astronomer/oss-integrations-private/issues/240

[Copilot]

Pull Request Overview

This PR fixes the behavior of ExecutionMode.WATCHER when used with DbtTaskGroup by refactoring how the producer watcher task establishes dependencies with root dbt nodes.

• Renamed _add_producer_watcher() to _add_producer_watcher_and_dependencies() to better reflect its expanded responsibilities
• Moved producer watcher creation to occur after task dependencies are established, enabling proper detection of root nodes
• Added workaround for Airflow bug that prevents setting producer task as upstream dependency in TaskGroups

Reviewed Changes

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

File	Description
cosmos/airflow/graph.py	Refactored producer watcher creation logic to properly handle dependencies for both DbtDag and DbtTaskGroup
tests/operators/test_watcher.py	Added new integration test for DbtTaskGroup with watcher mode and updated existing test documentation
dags/example_watcher.py	Created new example DAG demonstrating watcher mode usage
dev/dags/example_watcher.py	Removed documentation comment markers

---

_{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 (a9fca6a) to head (499ab48).
⚠️ Report is 2 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2044   +/-   ##
=======================================
  Coverage   97.70%   97.71%           
=======================================
  Files          89       89           
  Lines        5673     5681    +8     
=======================================
+ Hits         5543     5551    +8     
  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.

[pankajkoti]

Happy to merge if we've validated the behaviour. Don't have time myself today to validate the issue and the fix.