Push producer's task completion status to XCOM

cosmos/operators/watcher.py

@@ -97,32 +97,40 @@ def _finalize(self, context: Context, startup_events: list[dict[str, Any]]) -> N
             ti.xcom_push(key="dbt_startup_events", value=startup_events)
 
     def execute(self, context: Context, **kwargs: Any) -> Any:
-        if not self.invocation_mode:
-            self._discover_invocation_mode()
+        try:
+            if not self.invocation_mode:
+                self._discover_invocation_mode()
 
-        use_events = self.invocation_mode == InvocationMode.DBT_RUNNER and EventMsg is not None
-        self.log.debug("DbtProducerWatcherOperator: use_events=%s", use_events)
+            use_events = self.invocation_mode == InvocationMode.DBT_RUNNER and EventMsg is not None
+            self.log.debug("DbtProducerWatcherOperator: use_events=%s", use_events)
 
-        startup_events: list[dict[str, Any]] = []
+            startup_events: list[dict[str, Any]] = []
 
-        if use_events:
+            if use_events:
+
+                def _callback(ev: EventMsg) -> None:
+                    name = ev.info.name
+                    if name in {"MainReportVersion", "AdapterRegistered"}:
+                        self._handle_startup_event(ev, startup_events)
+                    elif name == "NodeFinished":
+                        self._handle_node_finished(ev, context)
 
-            def _callback(ev: EventMsg) -> None:
-                name = ev.info.name
-                if name in {"MainReportVersion", "AdapterRegistered"}:
-                    self._handle_startup_event(ev, startup_events)
-                elif name == "NodeFinished":
-                    self._handle_node_finished(ev, context)
+                self._dbt_runner_callbacks = [_callback]
+                result = super().execute(context=context, **kwargs)
 
-            self._dbt_runner_callbacks = [_callback]
-            result = super().execute(context=context, **kwargs)
+                self._finalize(context, startup_events)
+                return_value = result
+            else:
+                # Fallback – push run_results.json via base class helper
+                kwargs["push_run_results_to_xcom"] = True
+                return_value = super().execute(context=context, **kwargs)
 
-            self._finalize(context, startup_events)
-            return result
+            context["ti"].xcom_push(key="task_status", value="completed")
+            return return_value
 
-        # Fallback – push run_results.json via base class helper
-        kwargs["push_run_results_to_xcom"] = True
-        return super().execute(context=context, **kwargs)
+        except Exception:
+            context["ti"].xcom_push(key="task_status", value="completed")
+            raise
 
 
 class DbtConsumerWatcherSensor(BaseSensorOperator, DbtRunLocalOperator):  # type: ignore[misc]

tests/operators/test_watcher.py

@@ -70,6 +70,40 @@ def test_dbt_producer_watcher_operator_priority_weight_override():
     assert op.priority_weight == 100
 
 
+def test_dbt_producer_watcher_operator_pushes_completion_status():
+    """Test that operator pushes 'completed' status to XCom in both success and failure cases."""
+    op = DbtProducerWatcherOperator(project_dir=".", profile_config=None)
+    mock_ti = _MockTI()
+    context = {"ti": mock_ti}
+
+    # Test success case
+    with patch("cosmos.operators.local.DbtLocalBaseOperator.execute") as mock_execute:
+        op.execute(context=context)
+
+        # Verify status was pushed
+        assert mock_ti.store.get("task_status") == "completed"
+        # Verify parent execute was called
+        mock_execute.assert_called_once()
+
+    # Reset mock and store
+    mock_ti.store.clear()
+
+    # Test failure case
+    class TestException(Exception):
+        pass
+
+    with patch("cosmos.operators.local.DbtLocalBaseOperator.execute") as mock_execute:
+        mock_execute.side_effect = TestException("test error")
+
+        with pytest.raises(TestException):
+            op.execute(context=context)
+
+        # Verify completed status was pushed even in failure case
+        assert mock_ti.store.get("task_status") == "completed"
+        # Verify parent execute was called
+        mock_execute.assert_called_once()
+
+
 def test_handle_startup_event():
     op = DbtProducerWatcherOperator(project_dir=".", profile_config=None)
     lst: list[dict] = []