raghavan-arvind commented on code in PR #71625:
URL: https://github.com/apache/airflow/pull/71625#discussion_r3785513236
##########
providers/datadog/src/airflow/providers/datadog/sensors/datadog.py:
##########
@@ -98,3 +119,80 @@ def poke(self, context: Context) -> bool:
# If no check was inserted, assume any event that matched yields true.
return bool(response)
+
+
+class DatadogMonitorSensor(BaseSensorOperator):
+ """
+ Waits for a Datadog monitor to reach one of the target states.
+
+ :param monitor_id: The id of the Datadog monitor to watch.
+ :param target_states: Monitor overall states that complete the wait
+ (e.g. ``("OK",)`` or ``("Alert", "Warn")``).
+ :param datadog_conn_id: The connection to datadog, containing metadata for
api keys.
+ :param deferrable: Run sensor in deferrable mode.
+ """
+
+ ui_color = "#66c3dd"
+
+ def __init__(
+ self,
+ *,
+ monitor_id: int,
+ target_states: Sequence[str] = ("OK",),
+ datadog_conn_id: str = "datadog_default",
+ deferrable: bool = conf.getboolean("operators", "default_deferrable",
fallback=False),
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.monitor_id = monitor_id
+ self.target_states = target_states
+ self.datadog_conn_id = datadog_conn_id
+ self.deferrable = deferrable
+
+ def poke(self, context: Context) -> bool:
+ state = get_monitor_state(self.monitor_id, self.datadog_conn_id)
+ self.log.info("Monitor %s overall_state=%s", self.monitor_id, state)
+ return state in self.target_states
+
+ def execute(self, context: Context) -> None:
+ if not self.deferrable:
+ super().execute(context)
+ return
+ if self.poke(context):
+ return
+ self.defer(
+ trigger=DatadogMonitorTrigger(
+ monitor_id=self.monitor_id,
+ target_states=self.target_states,
+ datadog_conn_id=self.datadog_conn_id,
+ poke_interval=self.poke_interval,
+ ),
+ method_name="execute_complete",
+ timeout=timedelta(seconds=self.timeout),
+ )
+
+ def execute_complete(self, context: Context, event: dict[str, Any]) ->
None:
+ if event.get("status") != "success":
+ raise AirflowException(f"DatadogMonitorTrigger failed: {event}")
+ self.log.info("Monitor %s reached state %s", self.monitor_id,
event.get("state"))
+
+ def resume_execution(self, next_method: str, next_kwargs: dict[str, Any] |
None, context: Context):
+ """
+ Resume from deferral, applying ``soft_fail`` only to timeouts.
+
+ ``BaseSensorOperator.resume_execution`` converts any deferral-path
failure into a skip when
+ ``soft_fail`` is set, including trigger crashes (e.g. a bad monitor id
or auth failure),
+ which silently skips the downstream branch. This sensor instead keeps
parity with poke and
+ reschedule modes, where only timeouts are skippable and other errors
fail the task.
+ """
+ try:
+ return super(BaseSensorOperator,
self).resume_execution(next_method, next_kwargs, context)
+ except TaskDeferralError as e:
+ timed_out = isinstance(e, TaskDeferralTimeout) or str(e) ==
"Trigger/execution timeout"
+ if timed_out and self.soft_fail:
+ raise AirflowSkipException(str(e)) from e
+ if getattr(self, "never_fail", False):
+ raise AirflowSkipException(str(e)) from e
+ if timed_out:
+ raise AirflowSensorTimeout(*e.args) from e
+ raise
Review Comment:
Ideally we can remove this and merge some variant of #71256 first.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
