diff --git a/airflow/serialization/serialized_objects.py b/airflow/serialization/serialized_objects.py index 4d6a5b7e57ecc..1513a50d38514 100644 --- a/airflow/serialization/serialized_objects.py +++ b/airflow/serialization/serialized_objects.py @@ -27,7 +27,7 @@ import weakref from inspect import signature from textwrap import dedent -from typing import TYPE_CHECKING, Any, Collection, Iterable, Mapping, NamedTuple, Union +from typing import TYPE_CHECKING, Any, Collection, Iterable, Mapping, NamedTuple, Union, cast import attrs import lazy_object_proxy @@ -350,6 +350,42 @@ def decode_priority_weight_strategy(var: str) -> PriorityWeightStrategy: return priority_weight_strategy_class() +def encode_start_trigger_args(var: StartTriggerArgs) -> dict[str, Any]: + """ + Encode a StartTriggerArgs. + + :meta private: + """ + serialize_kwargs = ( + lambda key: BaseSerialization.serialize(getattr(var, key)) if getattr(var, key) is not None else None + ) + return { + "__type": "START_TRIGGER_ARGS", + "trigger_cls": var.trigger_cls, + "trigger_kwargs": serialize_kwargs("trigger_kwargs"), + "next_method": var.next_method, + "next_kwargs": serialize_kwargs("next_kwargs"), + "timeout": var.timeout.total_seconds() if var.timeout else None, + } + + +def decode_start_trigger_args(var: dict[str, Any]) -> StartTriggerArgs: + """ + Decode a StartTriggerArgs. + + :meta private: + """ + deserialize_kwargs = lambda key: BaseSerialization.deserialize(var[key]) if var[key] is not None else None + + return StartTriggerArgs( + trigger_cls=var["trigger_cls"], + trigger_kwargs=deserialize_kwargs("trigger_kwargs"), + next_method=var["next_method"], + next_kwargs=deserialize_kwargs("next_kwargs"), + timeout=datetime.timedelta(seconds=var["timeout"]) if var["timeout"] else None, + ) + + class _XComRef(NamedTuple): """ Store info needed to create XComArg. @@ -1088,7 +1124,7 @@ def _serialize_node(cls, op: BaseOperator | MappedOperator, include_deps: bool) serialize_op["_is_empty"] = op.inherits_from_empty_operator serialize_op["start_trigger_args"] = ( - op.start_trigger_args.serialize() if op.start_trigger_args else None + encode_start_trigger_args(op.start_trigger_args) if op.start_trigger_args else None ) serialize_op["start_from_trigger"] = op.start_from_trigger @@ -1276,8 +1312,10 @@ def populate_operator(cls, op: Operator, encoded_op: dict[str, Any]) -> None: setattr(op, "_is_empty", bool(encoded_op.get("_is_empty", False))) start_trigger_args = None - if encoded_op.get("start_trigger_args", None): - start_trigger_args = StartTriggerArgs(**encoded_op.get("start_trigger_args", None)) + encoded_start_trigger_args = encoded_op.get("start_trigger_args", None) + if encoded_start_trigger_args: + encoded_start_trigger_args = cast(dict, encoded_start_trigger_args) + start_trigger_args = decode_start_trigger_args(encoded_start_trigger_args) setattr(op, "start_trigger_args", start_trigger_args) setattr(op, "start_from_trigger", bool(encoded_op.get("start_from_trigger", False))) diff --git a/airflow/triggers/base.py b/airflow/triggers/base.py index 190e2983cee07..7b5338ad2fa19 100644 --- a/airflow/triggers/base.py +++ b/airflow/triggers/base.py @@ -47,15 +47,6 @@ class StartTriggerArgs: next_kwargs: dict[str, Any] | None = None timeout: timedelta | None = None - def serialize(self): - return { - "trigger_cls": self.trigger_cls, - "trigger_kwargs": self.trigger_kwargs, - "next_method": self.next_method, - "next_kwargs": self.next_kwargs, - "timeout": self.timeout, - } - class BaseTrigger(abc.ABC, LoggingMixin): """ diff --git a/tests/serialization/test_dag_serialization.py b/tests/serialization/test_dag_serialization.py index 8650378854a50..cefc5b7a2dcf3 100644 --- a/tests/serialization/test_dag_serialization.py +++ b/tests/serialization/test_dag_serialization.py @@ -2284,8 +2284,8 @@ def test_start_trigger_args_in_serialized_dag(self): class TestOperator(BaseOperator): start_trigger_args = StartTriggerArgs( - trigger_cls="airflow.triggers.testing.SuccessTrigger", - trigger_kwargs=None, + trigger_cls="airflow.triggers.temporal.TimeDeltaTrigger", + trigger_kwargs={"delta": timedelta(seconds=1)}, next_method="execute_complete", next_kwargs=None, timeout=None, @@ -2294,7 +2294,7 @@ class TestOperator(BaseOperator): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self.start_trigger_args.trigger_kwargs = {} + self.start_trigger_args.trigger_kwargs = {"delta": timedelta(seconds=2)} self.start_from_trigger = True def execute_complete(self): @@ -2323,16 +2323,28 @@ def execute_complete(self): Test2Operator(task_id="test_task_2") serialized_obj = SerializedDAG.to_dict(dag) - - for task in serialized_obj["dag"]["tasks"]: - assert task["__var"]["start_trigger_args"] == { - "trigger_cls": "airflow.triggers.testing.SuccessTrigger", - "trigger_kwargs": {}, - "next_method": "execute_complete", - "next_kwargs": None, - "timeout": None, - } - assert task["__var"]["start_from_trigger"] is True + tasks = serialized_obj["dag"]["tasks"] + + assert tasks[0]["__var"]["start_trigger_args"] == { + "__type": "START_TRIGGER_ARGS", + "trigger_cls": "airflow.triggers.temporal.TimeDeltaTrigger", + # "trigger_kwargs": {"__type": "dict", "__var": {"delta": {"__type": "timedelta", "__var": 2.0}}}, + "trigger_kwargs": {"__type": "dict", "__var": {"delta": {"__type": "timedelta", "__var": 2.0}}}, + "next_method": "execute_complete", + "next_kwargs": None, + "timeout": None, + } + assert tasks[0]["__var"]["start_from_trigger"] is True + + assert tasks[1]["__var"]["start_trigger_args"] == { + "__type": "START_TRIGGER_ARGS", + "trigger_cls": "airflow.triggers.testing.SuccessTrigger", + "trigger_kwargs": {"__type": "dict", "__var": {}}, + "next_method": "execute_complete", + "next_kwargs": None, + "timeout": None, + } + assert tasks[1]["__var"]["start_from_trigger"] is True def test_kubernetes_optional():