Watcher skip on producer failures

cosmos/airflow/graph.py

@@ -6,12 +6,15 @@
 from typing import Any
 
 try:  # Airflow 3
+    from airflow.providers.standard.operators.empty import EmptyOperator
     from airflow.sdk.bases.operator import BaseOperator
 except ImportError:  # Airflow 2
     from airflow.models import BaseOperator
+    from airflow.operators.empty import EmptyOperator  # type: ignore[no-redef]
 
 from airflow.models.base import ID_LEN as AIRFLOW_MAX_ID_LENGTH
 from airflow.models.dag import DAG
+from airflow.utils.trigger_rule import TriggerRule
 
 try:
     # Airflow 3.1 onwards
@@ -679,7 +682,7 @@ def _add_watcher_producer_task(
     render_config: RenderConfig | None = None,
     execution_mode: ExecutionMode = ExecutionMode.WATCHER,
     tests_per_model: dict[str, list[str]] | None = None,
-) -> BaseOperator:
+) -> tuple[BaseOperator, EmptyOperator]:
     """
     Create the producer task for the watcher execution mode and add it to the tasks_map.
     The producer task is the task that will be used to produce the events for the watcher execution mode.
@@ -711,12 +714,22 @@ def _add_watcher_producer_task(
     )
     producer_airflow_task = create_airflow_task(producer_task_metadata, dag, task_group=task_group)
     tasks_map[PRODUCER_WATCHER_TASK_ID] = producer_airflow_task
-    return producer_airflow_task
+
+    producer_task_gate = EmptyOperator(  # type: ignore[no-untyped-call]
+        task_id=f"{PRODUCER_WATCHER_TASK_ID}_gate",
+        dag=dag,
+        task_group=task_group,
+        trigger_rule=TriggerRule.NONE_FAILED,
+        depends_on_past=producer_airflow_task.depends_on_past,
+    )
+    producer_airflow_task >> producer_task_gate
+    return producer_airflow_task, producer_task_gate
 
 
 def _add_watcher_dependencies(
     dag: DAG,
     producer_airflow_task: BaseOperator,
+    producer_gate: BaseOperator,
     task_args: dict[str, Any],
     tasks_map: dict[str, Any],
     nodes: dict[str, DbtNode] | None = None,
@@ -728,7 +741,7 @@ def _add_watcher_dependencies(
     """
     for node_id, task_or_taskgroup in tasks_map.items():
         # We do not want to set a dependency between the producer task and itself
-        if node_id == PRODUCER_WATCHER_TASK_ID:
+        if node_id == PRODUCER_WATCHER_TASK_ID or node_id == f"{PRODUCER_WATCHER_TASK_ID}_gate":
             continue
 
         node_tasks = (
@@ -758,6 +771,10 @@ def _add_watcher_dependencies(
                 for task in always_run_tasks:
                     task.trigger_rule = task_args.get("trigger_rule", "always")  # type: ignore[attr-defined]
 
+        # If depends_on_past isn't true then gating all the tasks isn't really needed.
+        if producer_airflow_task.wait_for_downstream and not task_or_taskgroup.downstream_task_ids:
+            task_or_taskgroup >> producer_gate
+
 
 def should_create_detached_nodes(render_config: RenderConfig) -> bool:
     """
@@ -908,7 +925,7 @@ def build_airflow_graph(  # noqa: C901 TODO: https://github.com/astronomer/astro
     task_groups: dict[str, TaskGroup] = {}
     task_or_group: TaskGroup | BaseOperator | None
     parent_task_group = task_group
-    producer_task: BaseOperator | None = None
+    producer_tasks: tuple[BaseOperator, EmptyOperator] | None = None
 
     # Identify test nodes that should be run detached from the associated dbt resource nodes because they
     # have multiple parents
@@ -926,7 +943,7 @@ def build_airflow_graph(  # noqa: C901 TODO: https://github.com/astronomer/astro
         # We are intentionally creating the producer task ahead of the consumer tasks
         # Airflow priority weight is not being respected in multiple versions of the library, including 3.1
         # To instantiate the producer before helps having it before on the DAG topological order and scheduling this task before the consumer tasks
-        producer_task = _add_watcher_producer_task(
+        producer_tasks = _add_watcher_producer_task(
             dag=dag,
             task_args={**task_args, **setup_operator_args},
             tasks_map=tasks_map,
@@ -988,6 +1005,9 @@ def build_airflow_graph(  # noqa: C901 TODO: https://github.com/astronomer/astro
         leaves_ids = calculate_leaves(tasks_ids=list(tasks_map.keys()), nodes=nodes)
         for leaf_node_id in leaves_ids:
             tasks_map[leaf_node_id] >> test_task
+        if producer_tasks and producer_tasks[0].depends_on_past:
+            test_task >> producer_tasks[1]
+            test_task.wait_for_downstream = True
     elif render_config.test_behavior in (TestBehavior.BUILD, TestBehavior.AFTER_EACH):
         # Handle detached test nodes
         for node_id, node in detached_nodes.items():
@@ -1012,10 +1032,11 @@ def build_airflow_graph(  # noqa: C901 TODO: https://github.com/astronomer/astro
 
     create_airflow_task_dependencies(nodes, tasks_map)
 
-    if producer_task:
+    if producer_tasks:
         _add_watcher_dependencies(
             dag=dag,
-            producer_airflow_task=producer_task,
+            producer_airflow_task=producer_tasks[0],
+            producer_gate=producer_tasks[1],
             task_args=task_args,
             tasks_map=tasks_map,
             nodes=nodes,

cosmos/operators/_watcher/base.py

@@ -288,6 +288,9 @@ def __init__(
         self.deferrable = deferrable
         self.model_unique_id = extra_context.get("dbt_node_config", {}).get("unique_id")
 
+        if self.depends_on_past:
+            self.wait_for_downstream = True
+
     @staticmethod
     def _filter_flags(flags: list[str]) -> list[str]:
         """Filters out dbt flags that are incompatible with retry (e.g., --select, --exclude)."""
@@ -502,11 +505,10 @@ def poke(self, context: Context) -> bool:
         _log_dbt_event(dbt_events)
 
         if status is None:
-
-            if producer_task_state == "failed":
+            if producer_task_state == "failed" or producer_task_state == "skipped":
                 if self.poke_retry_number > 0:
                     raise AirflowException(
-                        f"The dbt build command failed in producer task. Please check the log of task {self.producer_task_id} for details."
+                        f"The dbt build command {producer_task_state} in the producer task. Please check the log of task {self.producer_task_id} for details."
                     )
                 else:
                     # This handles the scenario of tasks that failed with `State.UPSTREAM_FAILED`

cosmos/operators/_watcher/triggerer.py

@@ -26,7 +26,6 @@
 
 
 class WatcherTrigger(BaseTrigger):
-
     def __init__(
         self,
         model_unique_id: str,
@@ -213,10 +212,11 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
                     event_data["compiled_sql"] = compiled_sql
                 yield TriggerEvent(event_data)  # type: ignore[no-untyped-call]
                 return
-            elif producer_task_state == "failed":
+            elif producer_task_state == "failed" or producer_task_state == "skipped":
                 logger.error(
-                    "Watcher producer task '%s' failed before delivering results for node '%s'",
+                    "Watcher producer task '%s' %s before delivering results for node '%s'",
                     self.producer_task_id,
+                    producer_task_state,
                     self.model_unique_id,
                 )
                 yield TriggerEvent({"status": EventStatus.FAILED, "reason": "producer_failed"})  # type: ignore[no-untyped-call]

cosmos/operators/watcher.py

@@ -8,7 +8,7 @@
 from datetime import timedelta
 from typing import TYPE_CHECKING, Any
 
-from airflow.exceptions import AirflowException
+from airflow.exceptions import AirflowException, AirflowSkipException
 
 from cosmos.config import ProfileConfig
 from cosmos.operators._watcher import _parse_compressed_xcom, safe_xcom_push
@@ -120,6 +120,9 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         if self.invocation_mode == InvocationMode.SUBPROCESS:
             self.log_format = "json"
 
+        if self.depends_on_past:
+            self.wait_for_downstream = True
+
     @staticmethod
     def _serialize_event(event_message: EventMsg) -> dict[str, Any]:
         """Convert structured dbt EventMsg to plain dict."""
@@ -187,12 +190,10 @@ def execute(self, context: Context, **kwargs: Any) -> Any:
         try_number = getattr(task_instance, "try_number", 1)
 
         if try_number > 1:
-            self.log.info(
-                "Dbt WATCHER producer task does not support Airflow retries. "
-                "Detected attempt #%s; skipping execution to avoid running a second dbt build.",
-                try_number,
+            raise AirflowSkipException(
+                "DbtProducerWatcherOperator does not support Airflow retries. "
+                f"Detected attempt #{try_number}; skipping execution to avoid running a second dbt build."
             )
-            return None
 
         self.log.info(
             "Dbt WATCHER producer task forces Airflow retries to 0 so the dbt build only runs once; "
@@ -238,9 +239,10 @@ def _callback(event_message: EventMsg) -> None:
             safe_xcom_push(task_instance=context["ti"], key="task_status", value="completed")
             return return_value
 
-        except Exception:
+        except Exception as e:
             safe_xcom_push(task_instance=context["ti"], key="task_status", value="completed")
-            raise
+            self.log.exception("DbtProducerWatcherOperator execution failed")
+            raise AirflowSkipException("Skipping execution due to task failure") from e
 
 
 class DbtConsumerWatcherSensor(BaseConsumerSensor, DbtRunLocalOperator):  # type: ignore[misc]
@@ -352,6 +354,8 @@ class DbtTestWatcherOperator(EmptyOperator):
     """
 
     def __init__(self, *args: Any, **kwargs: Any):
+        default_args = kwargs.get("default_args", {})
         desired_keys = ("dag", "task_group", "task_id")
         new_kwargs = {key: value for key, value in kwargs.items() if key in desired_keys}
-        super().__init__(**new_kwargs)  # type: ignore[no-untyped-call]
+        depends_on_past = kwargs.get("depends_on_past", False) or default_args.get("depends_on_past", False)
+        super().__init__(depends_on_past=depends_on_past, wait_for_downstream=depends_on_past, **new_kwargs)  # type: ignore[no-untyped-call]

cosmos/operators/watcher_kubernetes.py

@@ -13,7 +13,7 @@
         from airflow.utils.context import Context  # type: ignore[attr-defined]
 
 import kubernetes.client as k8s
-from airflow.exceptions import AirflowException
+from airflow.exceptions import AirflowException, AirflowSkipException, TaskDeferred
 from airflow.providers.cncf.kubernetes.callbacks import KubernetesPodOperatorCallback, client_type
 
 try:
@@ -44,7 +44,6 @@
 
 
 class WatcherKubernetesCallback(KubernetesPodOperatorCallback):  # type: ignore[misc]
-
     @staticmethod
     def progress_callback(
         *,
@@ -74,7 +73,6 @@ def progress_callback(
 
 
 class DbtProducerWatcherKubernetesOperator(DbtBuildKubernetesOperator):
-
     template_fields: tuple[str, ...] = tuple(DbtBuildKubernetesOperator.template_fields) + ("deferrable",)
     _process_log_line_callable: Callable[[str, dict[str, Any]], None] | None = store_dbt_resource_status_from_log
 
@@ -99,6 +97,9 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(task_id=task_id, *args, **kwargs)
         self.dbt_cmd_flags += ["--log-format", "json"]
 
+        if self.depends_on_past:
+            self.wait_for_downstream = True
+
     @cached_property
     def pod_manager(self) -> CosmosKubernetesPodManager:
         return CosmosKubernetesPodManager(kube_client=self.client, callbacks=self.callbacks)
@@ -113,18 +114,31 @@ def execute(self, context: Context, **kwargs: Any) -> Any:
         try_number = getattr(task_instance, "try_number", 1)
 
         if try_number > 1:
-            self.log.info(
+            raise AirflowSkipException(
                 "DbtProducerWatcherKubernetesOperator does not support Airflow retries. "
-                "Detected attempt #%s; skipping execution to avoid running a second dbt build.",
-                try_number,
+                f"Detected attempt #{try_number}; skipping execution to avoid running a second dbt build."
             )
-            return None
 
         # This global variable is used to make the task context available to the K8s callback.
         # While the callback is set during the operator initialization, the context is only created during the operator's execution.
         global producer_task_context
         producer_task_context = context
-        return super().execute(context, **kwargs)
+        try:
+            return super().execute(context, **kwargs)
+        except (AirflowSkipException, TaskDeferred):
+            raise
+        except Exception as e:
+            self.log.exception("Dbt execution failed")
+            raise AirflowSkipException("Skipping execution due to task failure") from e
+
+    def trigger_reentry(self, *args: Any, **kwargs: Any) -> Any:
+        try:
+            return super().trigger_reentry(*args, **kwargs)
+        except (AirflowSkipException, TaskDeferred):
+            raise
+        except Exception as e:
+            self.log.exception("Dbt execution failed")
+            raise AirflowSkipException("Skipping execution due to task failure") from e
 
 
 class DbtConsumerWatcherKubernetesSensor(BaseConsumerSensor, DbtRunKubernetesOperator):

tests/airflow/test_graph.py

@@ -107,7 +107,6 @@
     tags=["nightly"],
     config={"materialized": "table", "meta": {"cosmos": {"operator_kwargs": {"pool": "custom_pool"}}}},
 )
-
 sample_nodes_list = [parent_seed, parent_node, test_parent_node, child_node, child2_node]
 sample_nodes = {node.unique_id: node for node in sample_nodes_list}
 
@@ -1178,13 +1177,87 @@ def test_test_behavior_for_watcher_mode(test_behavior):
     tasks = dag.tasks
     if test_behavior == TestBehavior.NONE:
         for task in tasks:
-            assert not isinstance(task, DbtTestWatcherOperator or DbtTestLocalOperator)
-        assert len(tasks) == 5
-    if test_behavior == TestBehavior.AFTER_EACH:
+            assert not isinstance(task, (DbtTestWatcherOperator, DbtTestLocalOperator))
         assert len(tasks) == 6
+    if test_behavior == TestBehavior.AFTER_EACH:
+        assert len(tasks) == 7
     if test_behavior == TestBehavior.AFTER_ALL:
         assert any(isinstance(task, DbtTestLocalOperator) for task in tasks)
-        assert len(tasks) == 6
+        assert len(tasks) == 7
+
+
+@pytest.mark.parametrize("depends_on_past", [False, True])
+@pytest.mark.parametrize("test_behavior", [TestBehavior.NONE, TestBehavior.AFTER_EACH, TestBehavior.AFTER_ALL])
+def test_watcher_dependency_wiring(test_behavior, depends_on_past):
+    with DAG("test-id", start_date=datetime(2022, 1, 1), default_args={"depends_on_past": depends_on_past}) as dag:
+        task_args = {
+            "project_dir": SAMPLE_PROJ_PATH,
+            "conn_id": "fake_conn",
+            "profile_config": ProfileConfig(
+                profile_name="default",
+                target_name="default",
+                profile_mapping=PostgresUserPasswordProfileMapping(
+                    conn_id="fake_conn",
+                    profile_args={"schema": "public"},
+                ),
+            ),
+        }
+
+    child_2b = DbtNode(
+        unique_id=f"{DbtResourceType.MODEL.value}.{SAMPLE_PROJ_PATH.stem}.child2.v2_b",
+        resource_type=DbtResourceType.MODEL,
+        depends_on=[parent_node.unique_id],
+        path_base=SAMPLE_PROJ_PATH,
+        original_file_path=Path("gen3/models/child2_v2.sql"),
+        tags=["nightly"],
+        config={"materialized": "table", "meta": {"cosmos": {"operator_kwargs": {"pool": "custom_pool"}}}},
+        has_test=True,
+        has_non_detached_test=True,
+    )
+    child_2b_test = DbtNode(
+        unique_id=f"{DbtResourceType.TEST.value}.{SAMPLE_PROJ_PATH.stem}.child2.test_v2_b",
+        resource_type=DbtResourceType.TEST,
+        depends_on=[child_2b.unique_id],
+        path_base=Path("."),
+        original_file_path=Path("."),
+    )
+
+    build_airflow_graph(
+        nodes={child_2b.unique_id: child_2b, child_2b_test.unique_id: child_2b_test, **sample_nodes},
+        dag=dag,
+        execution_mode=ExecutionMode.WATCHER,
+        test_indirect_selection=TestIndirectSelection.EAGER,
+        task_args=task_args,
+        render_config=RenderConfig(
+            test_behavior=test_behavior,
+        ),
+        dbt_project_name="astro_shop",
+    )
+    if not depends_on_past:
+        assert dag.task_dict["dbt_producer_watcher_gate"].upstream_task_ids == {"dbt_producer_watcher"}
+        assert all(task.wait_for_downstream is False for task in dag.tasks)
+        return
+
+    assert all(task.wait_for_downstream is True for task in dag.tasks if task.task_id != "dbt_producer_watcher_gate")
+    if test_behavior == TestBehavior.NONE:
+        assert dag.task_dict["dbt_producer_watcher_gate"].upstream_task_ids == {
+            "child_run",
+            "dbt_producer_watcher",
+            "child2_v2_run",
+            "child2_v2_b_run",
+        }
+    if test_behavior == TestBehavior.AFTER_EACH:
+        assert dag.task_dict["dbt_producer_watcher_gate"].upstream_task_ids == {
+            "child_run",
+            "dbt_producer_watcher",
+            "child2_v2_run",
+            "child2_v2_b.test",
+        }
+    if test_behavior == TestBehavior.AFTER_ALL:
+        assert dag.task_dict["dbt_producer_watcher_gate"].upstream_task_ids == {
+            "dbt_producer_watcher",
+            "astro_shop_test",
+        }
 
 
 def test_custom_meta():