Add validate tirgger events in dbt cloud deferrable tasks

generated/known_airflow_exceptions.txt

@@ -190,7 +190,6 @@ providers/databricks/tests/unit/databricks/hooks/test_databricks_base.py::1
 providers/datadog/src/airflow/providers/datadog/hooks/datadog.py::2
 providers/datadog/src/airflow/providers/datadog/sensors/datadog.py::1
 providers/dbt/cloud/src/airflow/providers/dbt/cloud/hooks/dbt.py::3
-providers/dbt/cloud/src/airflow/providers/dbt/cloud/sensors/dbt.py::1
 providers/dingding/src/airflow/providers/dingding/hooks/dingding.py::2
 providers/discord/src/airflow/providers/discord/operators/discord_webhook.py::1
 providers/docker/src/airflow/providers/docker/decorators/docker.py::1

providers/dbt/cloud/src/airflow/providers/dbt/cloud/hooks/dbt.py

@@ -138,6 +138,32 @@ class DbtCloudResourceLookupError(AirflowException):
     """Exception raised when a dbt Cloud resource cannot be uniquely identified."""
 
 
+class DbtCloudTriggerEventException(AirflowException):
+    """Raised when a deferred task resumes with a missing or malformed trigger event."""
+
+
+#: Statuses the provider's trigger emits in its terminal event.
+TRIGGER_EVENT_STATUSES = frozenset({"success", "cancelled", "error", "timeout"})
+
+
+def validate_execute_complete_event(event: dict[str, Any] | None = None) -> dict[str, Any]:
+    """
+    Validate the event a deferred task resumes with, returning it if well-formed.
+
+    The event crosses the triggerer/worker boundary through the metadata DB, so a
+    resuming task can receive ``None`` (e.g. a lost payload) or a status its handlers
+    do not recognize (triggerer/worker version skew, a custom trigger). Both must fail
+    loudly rather than silently falling through to the success path.
+    """
+    if event is None:
+        raise DbtCloudTriggerEventException("Trigger error: event is None")
+    if event.get("status") not in TRIGGER_EVENT_STATUSES:
+        raise DbtCloudTriggerEventException(
+            f"Unexpected trigger event status {event.get('status')!r}: {event!r}"
+        )
+    return event
+
+
 T = TypeVar("T", bound=Any)
 
 

providers/dbt/cloud/src/airflow/providers/dbt/cloud/operators/dbt.py

@@ -29,6 +29,7 @@
     DbtCloudJobRunException,
     DbtCloudJobRunStatus,
     JobRunInfo,
+    validate_execute_complete_event,
 )
 from airflow.providers.dbt.cloud.triggers.dbt import DbtCloudRunJobTrigger
 from airflow.providers.dbt.cloud.utils.openlineage import generate_openlineage_events_from_dbt_cloud_run
@@ -283,8 +284,9 @@ def execute(self, context: Context):
                 )
             return self.run_id
 
-    def execute_complete(self, context: Context, event: dict[str, Any]) -> int:
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> int:
         """Execute when the trigger fires - returns immediately."""
+        event = validate_execute_complete_event(event)
         self.run_id = event["run_id"]
         if event["status"] == "cancelled":
             raise DbtCloudJobRunException(f"Job run {self.run_id} has been cancelled.")

providers/dbt/cloud/src/airflow/providers/dbt/cloud/sensors/dbt.py

@@ -20,8 +20,13 @@
 from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
-from airflow.providers.common.compat.sdk import AirflowException, BaseSensorOperator, conf
-from airflow.providers.dbt.cloud.hooks.dbt import DbtCloudHook, DbtCloudJobRunException, DbtCloudJobRunStatus
+from airflow.providers.common.compat.sdk import BaseSensorOperator, conf
+from airflow.providers.dbt.cloud.hooks.dbt import (
+    DbtCloudHook,
+    DbtCloudJobRunException,
+    DbtCloudJobRunStatus,
+    validate_execute_complete_event,
+)
 from airflow.providers.dbt.cloud.triggers.dbt import DbtCloudRunJobTrigger
 from airflow.providers.dbt.cloud.utils.openlineage import generate_openlineage_events_from_dbt_cloud_run
 
@@ -114,15 +119,16 @@ def execute(self, context: Context) -> None:
                     method_name="execute_complete",
                 )
 
-    def execute_complete(self, context: Context, event: dict[str, Any]) -> int:
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> int:
         """
         Execute when the trigger fires - returns immediately.
 
         This relies on trigger to throw an exception, otherwise it assumes
         execution was successful.
         """
-        if event["status"] in ["error", "cancelled"]:
-            raise AirflowException()
+        event = validate_execute_complete_event(event)
+        if event["status"] != "success":
+            raise DbtCloudJobRunException(event["message"])
         self.log.info(event["message"])
         return int(event["run_id"])
 

providers/dbt/cloud/tests/unit/dbt/cloud/hooks/test_dbt.py

@@ -35,8 +35,10 @@
     DbtCloudJobRunException,
     DbtCloudJobRunStatus,
     DbtCloudResourceLookupError,
+    DbtCloudTriggerEventException,
     TokenAuth,
     fallback_to_default_account,
+    validate_execute_complete_event,
 )
 
 from tests_common.test_utils.compat import timezone
@@ -158,6 +160,34 @@ def test_invalid_terminal_job_run_status(self, statuses):
             DbtCloudJobRunStatus.check_is_valid(statuses)
 
 
+class TestValidateExecuteCompleteEvent:
+    @pytest.mark.parametrize(
+        ("event", "match"),
+        [
+            pytest.param(None, "event is None", id="none"),
+            pytest.param({}, "Unexpected trigger event status None", id="missing-status"),
+            pytest.param(
+                {"status": "ended", "run_id": 1234}, "Unexpected trigger event status", id="unknown-status"
+            ),
+        ],
+    )
+    def test_invalid_event_raises(self, event, match):
+        with pytest.raises(DbtCloudTriggerEventException, match=match):
+            validate_execute_complete_event(event)
+
+    @pytest.mark.parametrize(
+        "event",
+        [
+            pytest.param({"status": "success", "run_id": 1234, "message": "ok"}, id="success"),
+            pytest.param({"status": "cancelled", "run_id": 1234, "message": "cancelled"}, id="cancelled"),
+            pytest.param({"status": "error", "run_id": 1234, "message": "failed"}, id="error"),
+            pytest.param({"status": "timeout", "run_id": 1234, "message": "timed out"}, id="timeout"),
+        ],
+    )
+    def test_valid_event_is_returned(self, event):
+        assert validate_execute_complete_event(event) is event
+
+
 class TestDbtCloudHook:
     # TODO: Potential performance issue, converted setup_class to a setup_connections function level fixture
     @pytest.fixture(autouse=True)

providers/dbt/cloud/tests/unit/dbt/cloud/operators/test_dbt.py

@@ -24,7 +24,12 @@
 
 from airflow.models import DAG, Connection
 from airflow.providers.common.compat.sdk import TaskDeferred, timezone
-from airflow.providers.dbt.cloud.hooks.dbt import DbtCloudHook, DbtCloudJobRunException, DbtCloudJobRunStatus
+from airflow.providers.dbt.cloud.hooks.dbt import (
+    DbtCloudHook,
+    DbtCloudJobRunException,
+    DbtCloudJobRunStatus,
+    DbtCloudTriggerEventException,
+)
 from airflow.providers.dbt.cloud.operators.dbt import (
     DbtCloudGetJobRunArtifactOperator,
     DbtCloudListJobRunsOperator,
@@ -375,6 +380,26 @@ def test_execute_complete_timeout_cancel_job_does_not_mask_original_error(self):
             run_id=RUN_ID,
         )
 
+    @pytest.mark.parametrize(
+        "event",
+        [
+            pytest.param(None, id="none"),
+            pytest.param({"status": "ended", "run_id": RUN_ID, "message": "m"}, id="unknown-status"),
+        ],
+    )
+    def test_execute_complete_invalid_event_raises_instead_of_succeeding(self, event):
+        """Verify a missing or unrecognized trigger event fails loudly instead of succeeding."""
+        operator = DbtCloudRunJobOperator(
+            task_id=TASK_ID,
+            dbt_cloud_conn_id=ACCOUNT_ID_CONN,
+            job_id=JOB_ID,
+            dag=self.dag,
+            deferrable=True,
+        )
+
+        with pytest.raises(DbtCloudTriggerEventException):
+            operator.execute_complete(context=self.mock_context, event=event)
+
     @patch(
         "airflow.providers.dbt.cloud.hooks.dbt.DbtCloudHook.get_job_by_name",
         return_value=mock_response_json(DEFAULT_ACCOUNT_JOB_RESPONSE),

providers/dbt/cloud/tests/unit/dbt/cloud/sensors/test_dbt.py

@@ -22,8 +22,13 @@
 import pytest
 
 from airflow.models.connection import Connection
-from airflow.providers.common.compat.sdk import AirflowException, TaskDeferred
-from airflow.providers.dbt.cloud.hooks.dbt import DbtCloudHook, DbtCloudJobRunException, DbtCloudJobRunStatus
+from airflow.providers.common.compat.sdk import TaskDeferred
+from airflow.providers.dbt.cloud.hooks.dbt import (
+    DbtCloudHook,
+    DbtCloudJobRunException,
+    DbtCloudJobRunStatus,
+    DbtCloudTriggerEventException,
+)
 from airflow.providers.dbt.cloud.sensors.dbt import DbtCloudJobRunSensor
 from airflow.providers.dbt.cloud.triggers.dbt import DbtCloudRunJobTrigger
 
@@ -151,18 +156,42 @@ def test_execute_complete_success(self):
         [
             ("cancelled", "Job run 1234 has been cancelled."),
             ("error", "Job run 1234 has failed."),
+            ("timeout", "Job run 1234 has timed out."),
         ],
     )
     def test_execute_complete_failure(self, mock_status, mock_message):
-        """Assert execute_complete method to raise exception on the cancelled and error status"""
+        """Assert execute_complete method to raise exception on the cancelled, error, and timeout status"""
         task = DbtCloudJobRunSensor(
             dbt_cloud_conn_id=self.CONN_ID,
             task_id=self.TASK_ID,
             run_id=self.DBT_RUN_ID,
             timeout=self.TIMEOUT,
             deferrable=True,
         )
-        with pytest.raises(AirflowException):
+        with pytest.raises(DbtCloudJobRunException, match=mock_message):
             task.execute_complete(
                 context={}, event={"status": mock_status, "message": mock_message, "run_id": self.DBT_RUN_ID}
             )
+
+    @pytest.mark.parametrize(
+        ("event", "match"),
+        [
+            pytest.param(None, "event is None", id="none"),
+            pytest.param(
+                {"status": "rescheduling", "message": "m", "run_id": 1234},
+                "Unexpected trigger event status",
+                id="unknown-status",
+            ),
+        ],
+    )
+    def test_execute_complete_invalid_event_raises(self, event, match):
+        """Assert execute_complete raises instead of silently succeeding on a malformed event"""
+        task = DbtCloudJobRunSensor(
+            dbt_cloud_conn_id=self.CONN_ID,
+            task_id=self.TASK_ID,
+            run_id=self.DBT_RUN_ID,
+            timeout=self.TIMEOUT,
+            deferrable=True,
+        )
+        with pytest.raises(DbtCloudTriggerEventException, match=match):
+            task.execute_complete(context={}, event=event)