from collections.abc import Callable
from datetime import UTC, datetime, time, timedelta
from logging import getLogger
from typing import TYPE_CHECKING, Any, cast
from airflow_pydantic import fail, pass_
from airflow_pydantic.airflow import Param
if TYPE_CHECKING:
from airflow.providers.standard.operators.python import BranchPythonOperator, PythonOperator
from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator
from airflow.providers.standard.sensors.python import PythonSensor
else:
from airflow_pydantic.airflow import BranchPythonOperator, PythonOperator, PythonSensor, TriggerDagRunOperator
from .common import (
Action,
CheckResult,
Endtime,
FailTriggerKwargs,
MaxRetrigger,
PassTriggerKwargs,
PythonCallable,
ReferenceDate,
Result,
Runtime,
)
__all__ = (
"HighAvailabilityOperator",
"HighAvailabilitySensor",
)
_log = getLogger(__name__)
[docs]
class HighAvailabilitySensor(PythonSensor):
_original = "airflow_ha.operator.HighAvailabilitySensor"
_decide_task: BranchPythonOperator
_fail: PythonOperator
_retrigger_fail: TriggerDagRunOperator
_retrigger_pass: TriggerDagRunOperator
_stop_pass: PythonOperator
_stop_fail: PythonOperator
_pass_trigger_kwargs: PassTriggerKwargs | None = None
_pass_trigger_kwargs_conf: str = "{}"
_fail_trigger_kwargs: FailTriggerKwargs | None = None
_fail_trigger_kwargs_conf: str = "{}"
_check_end_conditions: Callable | None = None
_runtime: Runtime | None = None
_endtime: Endtime | None = None
_maxretrigger: MaxRetrigger | None = None
_reference_date: str | None = None
[docs]
def __init__(
self,
python_callable: PythonCallable,
pass_trigger_kwargs: PassTriggerKwargs | None = None,
fail_trigger_kwargs: FailTriggerKwargs | None = None,
*,
runtime: Runtime | None = None,
endtime: Endtime | None = None,
maxretrigger: MaxRetrigger | None = None,
reference_date: ReferenceDate = "data_interval_end",
**kwargs,
) -> None:
"""Run a health check and branch to stop or retrigger the DAG.
The callable returns a tuple of Result and Action enum members.
CONTINUE keeps the sensor polling. RETRIGGER creates another DAG run;
STOP finishes monitoring. Failed outcomes select a failure branch,
including after a failed retrigger schedules the next run.
Runtime and endtime limits select PASS/STOP. Exhausting maxretrigger
selects STOP while preserving the last health result.
"""
# These options control the behavior of the sensor
self._runtime = timedelta(seconds=runtime) if isinstance(runtime, int) else runtime
self._endtime = time.fromisoformat(endtime) if isinstance(endtime, str) else endtime
self._maxretrigger = maxretrigger or None
self._reference_date = reference_date or "data_interval_end"
# These are kwarsg to pass to the trigger operators
self._pass_trigger_kwargs = pass_trigger_kwargs or {}
self._fail_trigger_kwargs = fail_trigger_kwargs or {}
self._pass_trigger_kwargs_conf = self._pass_trigger_kwargs.pop("conf", {})
self._fail_trigger_kwargs_conf = self._fail_trigger_kwargs.pop("conf", {})
# Function to check end conditions
def _check_end_conditions_wrapper(
task_id=None,
runtime=self._runtime,
endtime=self._endtime,
maxretrigger=self._maxretrigger,
reference_date=self._reference_date,
**kwargs,
):
if task_id is None:
task_id = kwargs.get("task_id", self.task_id)
return _check_end_conditions(
task_id=task_id,
runtime=runtime,
endtime=endtime,
maxretrigger=maxretrigger,
reference_date=reference_date,
**kwargs,
)
self._check_end_conditions = _check_end_conditions_wrapper
# Function to control the sensor
callable_wrapper = lambda python_callable=python_callable, check_end_conditions=self._check_end_conditions, **kwargs: _callable_wrapper(
python_callable=python_callable, check_end_conditions=check_end_conditions, **kwargs
)
if not kwargs.get("trigger_rule"):
kwargs["trigger_rule"] = "none_failed"
callbacks = {
key: kwargs[key]
for key in ("on_failure_callback", "on_retry_callback", "on_execute_callback", "on_success_callback", "on_skipped_callback")
if key in kwargs
}
# Initialize the sensor
super().__init__(python_callable=callable_wrapper, **kwargs)
# Add params to dag to control overrides
self.dag.params.update(
{
f"{self.task_id}-force-run": Param(False, "Ignore runtime/endtime/maxretrigger and force run the task", type="boolean"),
f"{self.task_id}-force-runtime": Param(
None, "Override `runtime` in seconds, incompatible with `force-run`", type=["null", "integer"]
),
f"{self.task_id}-force-endtime": Param(
None, "Override `endtime`, incompatible with `force-run`", type=["null", "string"], format="time"
),
f"{self.task_id}-force-maxretrigger": Param(None, "Override `maxretrigger`, incompatible with `force-run`", type=["null", "integer"]),
f"{self.task_id}-force-retrigger": Param(None, "Override `retrigger`", type=["null", "integer"]),
}
)
# this is needed to ensure the dag fails, since the
# retrigger_fail step will pass (to ensure dag retriggers!)
self._fail = PythonOperator(
task_id=f"{self.task_id}-force-dag-fail", dag=self.dag, python_callable=fail, pool=kwargs.get("pool"), **callbacks
)
self._stop_pass = PythonOperator(
task_id=f"{self.task_id}-stop-pass", dag=self.dag, python_callable=pass_, pool=kwargs.get("pool"), **callbacks
)
self._stop_fail = PythonOperator(
task_id=f"{self.task_id}-stop-fail", dag=self.dag, python_callable=fail, pool=kwargs.get("pool"), **callbacks
)
# Update the retrigger counts in trigger kwargs
retrigger_count_conf = f'''{{{{ (dag_run.conf.get("{self.task_id}-retrigger", 0)|int) + 1 }}}}'''
referencedate_conf = f'''{{{{ dag_run.conf.get("{self.task_id}-referencedate", dag_run.start_date.isoformat()) }}}}'''
if isinstance(self._pass_trigger_kwargs_conf, dict):
self._pass_trigger_kwargs_conf[f"{self.task_id}-retrigger"] = retrigger_count_conf
self._pass_trigger_kwargs_conf[f"{self.task_id}-referencedate"] = referencedate_conf
else:
if not isinstance(self._pass_trigger_kwargs_conf, str) or not self._pass_trigger_kwargs_conf.strip().endswith("}"):
raise ValueError("pass_trigger_kwargs must be a dict or a JSON string")
self._pass_trigger_kwargs_conf = self._pass_trigger_kwargs_conf.strip()[:-1] + f', "{self.task_id}-retrigger": {retrigger_count_conf} }}'
self._pass_trigger_kwargs_conf = (
self._pass_trigger_kwargs_conf.strip()[:-1] + f', "{self.task_id}-referencedate": "{referencedate_conf}" }}'
)
if isinstance(self._fail_trigger_kwargs_conf, dict):
self._fail_trigger_kwargs_conf[f"{self.task_id}-retrigger"] = retrigger_count_conf
self._fail_trigger_kwargs_conf[f"{self.task_id}-referencedate"] = referencedate_conf
else:
if not isinstance(self._fail_trigger_kwargs_conf, str) or not self._fail_trigger_kwargs_conf.strip().endswith("}"):
raise ValueError("fail_trigger_kwargs must be a dict or a JSON string")
self._fail_trigger_kwargs_conf = self._fail_trigger_kwargs_conf.strip()[:-1] + f', "{self.task_id}-retrigger": {retrigger_count_conf} }}'
self._fail_trigger_kwargs_conf = (
self._fail_trigger_kwargs_conf.strip()[:-1] + f', "{self.task_id}-referencedate": "{referencedate_conf}" }}'
)
# Create the retrigger pass/fail operators
self._retrigger_fail = TriggerDagRunOperator(
task_id=f"{self.task_id}-retrigger-fail",
dag=self.dag,
conf=self._fail_trigger_kwargs_conf,
pool=kwargs.get("pool"),
**cast(dict[str, Any], {"trigger_dag_id": self.dag_id, "trigger_rule": "one_success", **callbacks, **self._fail_trigger_kwargs}),
)
self._retrigger_pass = TriggerDagRunOperator(
task_id=f"{self.task_id}-retrigger-pass",
dag=self.dag,
conf=self._pass_trigger_kwargs_conf,
pool=kwargs.get("pool"),
**cast(dict[str, Any], {"trigger_dag_id": self.dag_id, "trigger_rule": "one_success", **callbacks, **self._pass_trigger_kwargs}),
)
# Create the branch operator
branch_choices = {
(Result.PASS, Action.RETRIGGER): self._retrigger_pass.task_id,
(Result.PASS, Action.STOP): self._stop_pass.task_id,
(Result.FAIL, Action.RETRIGGER): self._retrigger_fail.task_id,
(Result.FAIL, Action.STOP): self._stop_fail.task_id,
}
choose_branch = lambda task_id=self.task_id, branch_choices=branch_choices, **kwargs: _choose_branch(
task_id=task_id, branch_choices=branch_choices, check_end_conditions=self._check_end_conditions, **kwargs
)
decide_task_args: dict[str, Any] = {
"task_id": f"{self.task_id}-decide",
"dag": self.dag,
"python_callable": choose_branch,
# NOTE: use none_skipped here as the sensor will fail in a timeout
"trigger_rule": "none_skipped",
"pool": kwargs.get("pool"),
**callbacks,
}
self._decide_task = BranchPythonOperator(**decide_task_args)
self >> self._decide_task
self._decide_task >> self._stop_pass
self._decide_task >> self._stop_fail
self._decide_task >> self._retrigger_pass
self._decide_task >> self._retrigger_fail >> self._fail
@property
def decide_task(self) -> PythonOperator:
return self._decide_task
@property
def stop_fail(self) -> PythonOperator:
return self._stop_fail
@property
def stop_pass(self) -> PythonOperator:
return self._stop_pass
@property
def retrigger_fail(self) -> TriggerDagRunOperator:
return self._retrigger_fail
@property
def retrigger_pass(self) -> TriggerDagRunOperator:
return self._retrigger_pass
@property
def check_end_conditions(self) -> Callable:
assert self._check_end_conditions is not None
return self._check_end_conditions
@property
def is_initial_run(self) -> Callable:
return lambda **kwargs: self.get_retrigger_count(**kwargs) == 0
@property
def get_retrigger_count(self) -> Callable:
return lambda **kwargs: int(kwargs["dag_run"].conf.get(f"{self.task_id}-retrigger", 0))
# Function to check end conditions
def _check_end_conditions(task_id, runtime, endtime, maxretrigger, reference_date, **kwargs):
# Check if force run in dag run kwargs
force_run_conf = kwargs["dag_run"].conf.get("airflow_ha_force_run", False)
force_run_param = kwargs["params"].get(f"{task_id}-force-run", False)
_log.info(f"airflow-ha configuration -- force_run (conf): {force_run_conf}, force_run (param): {force_run_param}")
# Grab the dag start date
dag_reference_date = kwargs["dag_run"].conf.get(f"{task_id}-referencedate")
if dag_reference_date is not None:
dag_reference_date = datetime.fromisoformat(dag_reference_date)
else:
dag_reference_date = getattr(kwargs["dag_run"], reference_date)
if not force_run_conf and not force_run_param:
runtime_override = kwargs["params"].get(f"{task_id}-force-runtime", None)
runtime = runtime if runtime_override is None else runtime_override
runtime = timedelta(seconds=runtime) if isinstance(runtime, int) else runtime
endtime = kwargs["params"].get(f"{task_id}-force-endtime", None) or endtime
endtime = time.fromisoformat(endtime) if isinstance(endtime, str) else endtime
# Check if runtime has exceeded
# NOTE: start date will always be normalize to UTC by airflow
elapsed_time = (datetime.now(tz=UTC) - dag_reference_date).total_seconds()
_log.info(
f"airflow-ha configuration -- runtime: {runtime}, reference_date({reference_date}): {dag_reference_date}, elapsed_time: {elapsed_time}"
)
if runtime is not None and elapsed_time > runtime.total_seconds():
# Runtime has exceeded, end
_log.info(f"runtime exceeded for {task_id}, stopping")
return Result.PASS, Action.STOP
# Check if endtime has passed
if endtime is not None:
# NOTE: we normalize to the DAG's timezone
# and use the provided endtime for convenience
dag = kwargs.get("dag")
assert dag is not None
dag_timezone = dag.timezone
endtime_as_datetime = datetime.combine(
date=dag_reference_date.astimezone(dag_timezone).date(),
time=endtime,
tzinfo=dag_timezone,
).astimezone(UTC)
_log.info(
f"airflow-ha configuration -- endtime: {endtime}, endtime_as_datetime: {endtime_as_datetime}, datetime.now(tz=UTC): {datetime.now(tz=UTC)}"
)
if endtime_as_datetime <= datetime.now(tz=UTC):
# Endtime has passed, end
_log.info(f"endtime passed for {task_id}, stopping")
return Result.PASS, Action.STOP
# Handle retrigger couts
maxretrigger = kwargs["params"].get(f"{task_id}-force-maxretrigger", None) or maxretrigger or -1
current_retrigger = int(kwargs["dag_run"].conf.get(f"{task_id}-retrigger", 0))
_log.info(f"airflow-ha configuration -- current_retrigger: {current_retrigger}, maxretrigger: {maxretrigger}")
# Check if maxretrigger has exceeded
if maxretrigger > -1 and current_retrigger >= maxretrigger:
# maxretrigger has exceeded, end
_log.info(f"maxretrigger exceeded for {task_id}: {maxretrigger} / {current_retrigger}, stopping")
return None, Action.STOP
# Function to control the sensor
def _callable_wrapper(python_callable, check_end_conditions, **kwargs):
task_instance = kwargs["task_instance"]
try:
ret: CheckResult = python_callable(**kwargs)
except Exception:
task_instance.xcom_push(key="return_value", value=(Result.FAIL, Action.RETRIGGER))
raise
if not isinstance(ret, tuple) or len(ret) != 2 or not isinstance(ret[0], Result) or not isinstance(ret[1], Action):
# malformed
task_instance.xcom_push(key="return_value", value=(Result.FAIL, Action.STOP))
return True
# push to xcom
task_instance.xcom_push(key="return_value", value=ret)
end_conditions = check_end_conditions(**kwargs)
if end_conditions is not None and end_conditions[1] == Action.STOP:
# end conditions met
_log.info(f"End conditions met, stopping: {end_conditions}")
return True
return ret[1] != Action.CONTINUE
def _choose_branch(branch_choices, task_id, check_end_conditions, **kwargs):
# Grab the task instance
task_instance = kwargs["task_instance"]
# Check if force run in dag run kwargs
end_conditions = check_end_conditions(task_id=task_id, **kwargs)
if end_conditions is not None and end_conditions[0] is not None:
return branch_choices[end_conditions]
if end_conditions is not None:
retrigger_exceeded = True
else:
retrigger_exceeded = False
# Otherwise, continue to evaluate
check_program_result = task_instance.xcom_pull(key="return_value", task_ids=task_id)
try:
if not isinstance(check_program_result, (tuple, list)) or len(check_program_result) != 2:
raise ValueError("Expected a result/action pair")
result = Result(check_program_result[0])
action = Action(check_program_result[1])
except (ValueError, TypeError):
_log.warning("Sensor returned no valid result, treating as failure")
result, action = Result.FAIL, Action.RETRIGGER
if retrigger_exceeded:
action = Action.STOP
elif action == Action.CONTINUE:
action = Action.RETRIGGER
ret = branch_choices[(result, action)]
_log.info(f"Sensor returned {result.name}, {action.name}, branching to {ret}")
return ret
# Alias
HighAvailabilityOperator = HighAvailabilitySensor