fix(providers/datadog): respect soft_fail argument when exception is …

…raised (#34472)

airflow/providers/datadog/sensors/datadog.py

@@ -21,7 +21,7 @@
 
 from datadog import api
 
-from airflow.exceptions import AirflowException
+from airflow.exceptions import AirflowException, AirflowSkipException
 from airflow.providers.datadog.hooks.datadog import DatadogHook
 from airflow.sensors.base import BaseSensorOperator
 
@@ -89,7 +89,11 @@ def poke(self, context: Context) -> bool:
 
         if isinstance(response, dict) and response.get("status", "ok") != "ok":
             self.log.error("Unexpected Datadog result: %s", response)
-            raise AirflowException("Datadog returned unexpected result")
+            # TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
+            message = "Datadog returned unexpected result"
+            if self.soft_fail:
+                raise AirflowSkipException(message)
+            raise AirflowException(message)
 
         if self.response_check:
             # run content check on response

tests/providers/datadog/sensors/test_datadog.py

@@ -20,6 +20,9 @@
 import json
 from unittest.mock import patch
 
+import pytest
+
+from airflow.exceptions import AirflowException, AirflowSkipException
 from airflow.models import Connection
 from airflow.providers.datadog.sensors.datadog import DatadogSensor
 from airflow.utils import db
@@ -111,3 +114,26 @@ def test_sensor_fail(self, api1, api2):
         )
 
         assert not sensor.poke({})
+
+    @pytest.mark.parametrize(
+        "soft_fail, expected_exception", ((False, AirflowException), (True, AirflowSkipException))
+    )
+    @patch("airflow.providers.datadog.hooks.datadog.api.Event.query")
+    @patch("airflow.providers.datadog.sensors.datadog.api.Event.query")
+    def test_sensor_fail_with_exception(self, api1, api2, soft_fail, expected_exception):
+        api1.return_value = zero_events
+        api2.return_value = {"status": "error"}
+
+        with pytest.raises(expected_exception):
+            sensor = DatadogSensor(
+                task_id="test_datadog",
+                datadog_conn_id="datadog_default",
+                from_seconds_ago=0,
+                up_to_seconds_from_now=0,
+                priority=None,
+                sources=None,
+                tags=None,
+                response_check=None,
+                soft_fail=soft_fail,
+            )
+            sensor.poke({})