diff --git a/task-sdk/src/airflow/sdk/bases/sensor.py b/task-sdk/src/airflow/sdk/bases/sensor.py index 0ca447189a983..66a97c10f0e4f 100644 --- a/task-sdk/src/airflow/sdk/bases/sensor.py +++ b/task-sdk/src/airflow/sdk/bases/sensor.py @@ -251,13 +251,18 @@ def run_duration() -> float: return xcom_value def resume_execution(self, next_method: str, next_kwargs: dict[str, Any] | None, context: Context): + # Use nested try/except to convert TaskDeferralTimeout to AirflowSensorTimeout + # while still allowing soft_fail/never_fail to handle both exception types. try: - return super().resume_execution(next_method, next_kwargs, context) - except TaskDeferralTimeout as e: - raise AirflowSensorTimeout(*e.args) from e + try: + return super().resume_execution(next_method, next_kwargs, context) + except TaskDeferralTimeout as e: + raise AirflowSensorTimeout(*e.args) from e except (AirflowException, TaskDeferralError) as e: if self.soft_fail: - raise AirflowSkipException(str(e)) from e + raise AirflowSkipException("Skipping due to soft_fail is set to True.") from e + if self.never_fail: + raise AirflowSkipException("Skipping due to never_fail is set to True.") from e raise def _get_next_poke_interval( diff --git a/task-sdk/tests/task_sdk/bases/test_sensor.py b/task-sdk/tests/task_sdk/bases/test_sensor.py index 8b9cc04a7860b..b37cdac3b80d4 100644 --- a/task-sdk/tests/task_sdk/bases/test_sensor.py +++ b/task-sdk/tests/task_sdk/bases/test_sensor.py @@ -31,6 +31,7 @@ AirflowSensorTimeout, AirflowSkipException, AirflowTaskTimeout, + TaskDeferralError, ) from airflow.models.trigger import TriggerFailureReason from airflow.providers.standard.operators.empty import EmptyOperator @@ -677,3 +678,59 @@ def test_fail_after_resuming_deferred_sensor(self, soft_fail, expected_exception async_sensor = DummyAsyncSensor(task_id="dummy_async_sensor", soft_fail=soft_fail) with pytest.raises(expected_exception): async_sensor.resume_execution("execute_complete", None, {}) + + @pytest.mark.parametrize( + ("soft_fail", "expected_exception"), + [ + (True, AirflowSkipException), + (False, AirflowSensorTimeout), + ], + ) + def test_timeout_after_resuming_deferred_sensor_with_soft_fail(self, soft_fail, expected_exception): + """Test that deferrable sensors with soft_fail skip on timeout instead of failing.""" + async_sensor = DummyAsyncSensor(task_id="dummy_async_sensor", soft_fail=soft_fail) + with pytest.raises(expected_exception): + async_sensor.resume_execution( + next_method="__fail__", + next_kwargs={"error": TriggerFailureReason.TRIGGER_TIMEOUT}, + context={}, + ) + + def test_timeout_after_resuming_deferred_sensor_with_never_fail(self): + """Test that deferrable sensors with never_fail skip on timeout.""" + async_sensor = DummyAsyncSensor(task_id="dummy_async_sensor", never_fail=True) + with pytest.raises(AirflowSkipException): + async_sensor.resume_execution( + next_method="__fail__", + next_kwargs={"error": TriggerFailureReason.TRIGGER_TIMEOUT}, + context={}, + ) + + @pytest.mark.parametrize( + ("soft_fail", "expected_exception"), + [ + (True, AirflowSkipException), + (False, TaskDeferralError), + ], + ) + def test_trigger_failure_after_resuming_deferred_sensor_with_soft_fail( + self, soft_fail, expected_exception + ): + """Test that deferrable sensors with soft_fail skip on trigger failure instead of failing.""" + async_sensor = DummyAsyncSensor(task_id="dummy_async_sensor", soft_fail=soft_fail) + with pytest.raises(expected_exception): + async_sensor.resume_execution( + next_method="__fail__", + next_kwargs={"error": TriggerFailureReason.TRIGGER_FAILURE}, + context={}, + ) + + def test_trigger_failure_after_resuming_deferred_sensor_with_never_fail(self): + """Test that deferrable sensors with never_fail skip on trigger failure.""" + async_sensor = DummyAsyncSensor(task_id="dummy_async_sensor", never_fail=True) + with pytest.raises(AirflowSkipException): + async_sensor.resume_execution( + next_method="__fail__", + next_kwargs={"error": TriggerFailureReason.TRIGGER_FAILURE}, + context={}, + )