diff --git a/airflow/models/baseoperator.py b/airflow/models/baseoperator.py index 458bd313bd42a..f13d0185e4938 100644 --- a/airflow/models/baseoperator.py +++ b/airflow/models/baseoperator.py @@ -973,23 +973,10 @@ def set_xcomargs_dependencies(self) -> None: """ from airflow.models.xcom_arg import XComArg - def apply_set_upstream(arg: Any): - if isinstance(arg, XComArg): - self.set_upstream(arg.operator) - elif isinstance(arg, (tuple, set, list)): - for elem in arg: - apply_set_upstream(elem) - elif isinstance(arg, dict): - for elem in arg.values(): - apply_set_upstream(elem) - elif hasattr(arg, "template_fields"): - for elem in arg.template_fields: - apply_set_upstream(elem) - for field in self.template_fields: if hasattr(self, field): arg = getattr(self, field) - apply_set_upstream(arg) + XComArg.apply_upstream_relationship(self, arg) @property def priority_weight_total(self) -> int: @@ -1734,6 +1721,8 @@ def __repr__(self) -> str: params: Union[ParamsDict, dict] = attr.ib(factory=ParamsDict) template_fields: Iterable[str] = attr.ib() + subdag: None = attr.ib(init=False) + @_is_dummy.default def _is_dummy_default(self): from airflow.operators.dummy import DummyOperator @@ -1795,12 +1784,17 @@ def from_decorator( return operator def __attrs_post_init__(self): + from airflow.models.xcom_arg import XComArg + if self.task_group: self.task_id = self.task_group.child_id(self.task_id) self.task_group.add(self) if self.dag: self.dag.add_task(self) + for arg in self.mapped_kwargs.values(): + XComArg.apply_upstream_relationship(self, arg) + @task_type.default def _default_task_type(self): # Can be a string if we are de-serialized @@ -1829,8 +1823,12 @@ def map(self, **kwargs) -> "MappedOperator": :return: ``self`` for easier method chaining """ + from airflow.models.xcom_arg import XComArg + if self.mapped_kwargs: raise RuntimeError("Already a mapped task") + for arg in kwargs.values(): + XComArg.apply_upstream_relationship(self, arg) return attr.evolve(self, mapped_kwargs=kwargs) @property @@ -1865,6 +1863,7 @@ def get_serialized_fields(cls): cls.__serialized_fields = frozenset( fields_dict.keys() - { + 'dag', 'deps', 'inherits_from_dummy_operator', 'operator_extra_links', diff --git a/airflow/models/xcom_arg.py b/airflow/models/xcom_arg.py index 5fe798b9ee3a1..772a3c8127fe9 100644 --- a/airflow/models/xcom_arg.py +++ b/airflow/models/xcom_arg.py @@ -17,12 +17,14 @@ from typing import TYPE_CHECKING, Any, List, Optional, Sequence, Union from airflow.exceptions import AirflowException -from airflow.models.baseoperator import BaseOperator, MappedOperator from airflow.models.taskmixin import DAGNode, DependencyMixin from airflow.models.xcom import XCOM_RETURN_KEY from airflow.utils.context import Context from airflow.utils.edgemodifier import EdgeModifier +if TYPE_CHECKING: + from airflow.models.baseoperator import BaseOperator, MappedOperator + class XComArg(DependencyMixin): """ @@ -59,9 +61,9 @@ class XComArg(DependencyMixin): :type key: str """ - def __init__(self, operator: Union[BaseOperator, MappedOperator], key: str = XCOM_RETURN_KEY): - self._operator = operator - self._key = key + def __init__(self, operator: "Union[BaseOperator, MappedOperator]", key: str = XCOM_RETURN_KEY): + self.operator = operator + self.key = key def __eq__(self, other): return self.operator == other.operator and self.key == other.key @@ -92,25 +94,15 @@ def __str__(self): xcom_pull = f"{{{{ task_instance.xcom_pull({xcom_pull_kwargs}) }}}}" return xcom_pull - @property - def operator(self) -> Union[BaseOperator, MappedOperator]: - """Returns operator of this XComArg.""" - return self._operator - @property def roots(self) -> List[DAGNode]: """Required by TaskMixin""" - return [self._operator] + return [self.operator] @property def leaves(self) -> List[DAGNode]: """Required by TaskMixin""" - return [self._operator] - - @property - def key(self) -> str: - """Returns keys of this XComArg""" - return self._key + return [self.operator] def set_upstream( self, @@ -144,3 +136,23 @@ def resolve(self, context: Context) -> Any: resolved_value = resolved_value[0] return resolved_value + + @staticmethod + def apply_upstream_relationship(op: "Union[BaseOperator, MappedOperator]", arg: Any): + """ + Set dependency for XComArgs. + + This looks for XComArg objects in ``arg`` "deeply" (looking inside lists, dicts and classes decorated + with "template_fields") and sets the relationship to ``op`` on any found. + """ + if isinstance(arg, XComArg): + op.set_upstream(arg.operator) + elif isinstance(arg, (tuple, set, list)): + for elem in arg: + XComArg.apply_upstream_relationship(op, elem) + elif isinstance(arg, dict): + for elem in arg.values(): + XComArg.apply_upstream_relationship(op, elem) + elif hasattr(arg, "template_fields"): + for elem in arg.template_fields: + XComArg.apply_upstream_relationship(op, elem) diff --git a/airflow/serialization/enums.py b/airflow/serialization/enums.py index a39baf4756753..f4227a6f7aed2 100644 --- a/airflow/serialization/enums.py +++ b/airflow/serialization/enums.py @@ -48,3 +48,4 @@ class DagAttributeTypes(str, Enum): TASK_GROUP = 'taskgroup' EDGE_INFO = 'edgeinfo' PARAM = 'param' + XCOM_REF = 'xcomref' diff --git a/airflow/serialization/serialized_objects.py b/airflow/serialization/serialized_objects.py index 785b9524adaa6..31d7aad7bea67 100644 --- a/airflow/serialization/serialized_objects.py +++ b/airflow/serialization/serialized_objects.py @@ -21,7 +21,7 @@ import logging from dataclasses import dataclass from inspect import Parameter, signature -from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Set, Type, Union +from typing import TYPE_CHECKING, Any, Dict, Iterable, List, NamedTuple, Optional, Set, Type, Union import cattr import pendulum @@ -36,6 +36,7 @@ from airflow.models.dag import DAG, create_timetable from airflow.models.param import Param, ParamsDict from airflow.models.taskmixin import DAGNode +from airflow.models.xcom_arg import XComArg from airflow.providers_manager import ProvidersManager from airflow.serialization.enums import DagAttributeTypes as DAT, Encoding from airflow.serialization.helpers import serialize_template_field @@ -168,6 +169,18 @@ def _decode_timetable(var: Dict[str, Any]) -> Timetable: return timetable_class.deserialize(var[Encoding.VAR]) +class _XcomRef(NamedTuple): + """ + Used to store info needed to create XComArg when deserializing MappedOperator. + + We can't turn it in to a XComArg until we've loaded _all_ the tasks, so when deserializing an operator we + need to create _something_, and then post-process it in deserialize_dag + """ + + task_id: str + key: str + + class BaseSerialization: """BaseSerialization provides utils for serialization.""" @@ -331,6 +344,8 @@ def _serialize(cls, var: Any) -> Any: # Unfortunately there is no support for r return SerializedTaskGroup.serialize_task_group(var) elif isinstance(var, Param): return cls._encode(cls._serialize_param(var), type_=DAT.PARAM) + elif isinstance(var, XComArg): + return cls._encode(cls._serialize_xcomarg(var), type_=DAT.XCOM_REF) else: log.debug('Cast type %s to str in serialization.', type(var)) return str(var) @@ -374,6 +389,8 @@ def _deserialize(cls, encoded_var: Any) -> Any: return tuple(cls._deserialize(v) for v in var) elif type_ == DAT.PARAM: return cls._deserialize_param(var) + elif type_ == DAT.XCOM_REF: + return cls._deserialize_xcomref(var) else: raise TypeError(f'Invalid type {type_!s} in deserialization.') @@ -476,6 +493,14 @@ def _deserialize_params_dict(cls, encoded_params: Dict) -> ParamsDict: return ParamsDict(op_params) + @classmethod + def _serialize_xcomarg(cls, arg: XComArg) -> dict: + return {"key": arg.key, "task_id": arg.operator.task_id} + + @classmethod + def _deserialize_xcomref(cls, encoded: dict) -> _XcomRef: + return _XcomRef(key=encoded['key'], task_id=encoded['task_id']) + class DependencyDetector: """Detects dependencies between DAGs.""" @@ -687,6 +712,8 @@ def deserialize_operator(cls, encoded_op: Dict[str, Any]) -> Union[BaseOperator, v = cls._deserialize_deps(v) elif k == "params": v = cls._deserialize_params_dict(v) + elif k in ("mapped_kwargs", "partial_kwargs"): + v = {arg: cls._deserialize(value) for arg, value in v.items()} elif k in cls._decorated_fields or k not in op.get_serialized_fields(): v = cls._deserialize(v) # else use v as it is @@ -970,6 +997,14 @@ def deserialize_dag(cls, encoded_dag: Dict[str, Any]) -> 'SerializedDAG': if serializable_task.subdag is not None: setattr(serializable_task.subdag, 'parent_dag', dag) + if isinstance(task, MappedOperator): + for d in (task.mapped_kwargs, task.partial_kwargs): + for k, v in d.items(): + if not isinstance(v, _XcomRef): + continue + + d[k] = XComArg(operator=dag.get_task(v.task_id), key=v.key) + for task_id in serializable_task.downstream_task_ids: # Bypass set_upstream etc here - it does more than we want dag.task_dict[task_id].upstream_task_ids.add(serializable_task.task_id) diff --git a/tests/models/test_baseoperator.py b/tests/models/test_baseoperator.py index 3e089c9acbd9c..92d273e3f6cc0 100644 --- a/tests/models/test_baseoperator.py +++ b/tests/models/test_baseoperator.py @@ -731,6 +731,21 @@ def test_map_unknown_arg_raises(): BaseOperator(task_id='a').map(file=[1, 2, {'a': 'b'}]) +def test_map_xcom_arg(): + """Test that dependencies are correct when mapping with an XComArg""" + from airflow.models.xcom_arg import XComArg + + with DAG("test-dag", start_date=DEFAULT_DATE): + task1 = BaseOperator(task_id="op1") + xcomarg = XComArg(task1, "test_key") + mapped = MockOperator(task_id='task_2').map(arg2=xcomarg) + finish = MockOperator(task_id="finish") + + mapped >> finish + + assert task1.downstream_list == [mapped] + + def test_partial_on_instance() -> None: """`.partial` on an instance should fail -- it's only designed to be called on classes""" with pytest.raises(TypeError): diff --git a/tests/serialization/test_dag_serialization.py b/tests/serialization/test_dag_serialization.py index acde989e3f682..0bf8b1e1bfe00 100644 --- a/tests/serialization/test_dag_serialization.py +++ b/tests/serialization/test_dag_serialization.py @@ -37,7 +37,7 @@ from airflow.hooks.base import BaseHook from airflow.kubernetes.pod_generator import PodGenerator from airflow.models import DAG, Connection, DagBag -from airflow.models.baseoperator import BaseOperator, BaseOperatorLink +from airflow.models.baseoperator import BaseOperator, BaseOperatorLink, MappedOperator from airflow.models.param import Param, ParamsDict from airflow.models.xcom import XCom from airflow.operators.bash import BashOperator @@ -52,7 +52,7 @@ from airflow.utils import timezone from airflow.utils.context import Context from airflow.utils.task_group import TaskGroup -from tests.test_utils.mock_operators import CustomOperator, CustomOpLink, GoogleLink +from tests.test_utils.mock_operators import CustomOperator, CustomOpLink, GoogleLink, MockOperator from tests.test_utils.timetables import CustomSerializationTimetable, cron_timetable, delta_timetable executor_config_pod = k8s.V1Pod( @@ -1573,7 +1573,10 @@ def mock__import__(name, globals_=None, locals_=None, fromlist=(), level=0): def test_mapped_operator_serde(): - real_op = BashOperator.partial(task_id='a').map(bash_command=[1, 2, {'a': 'b'}]) + literal = [1, 2, {'a': 'b'}] + real_op = BashOperator.partial(task_id='a', executor_config={'dict': {'sub': 'value'}}).map( + bash_command=literal + ) serialized = SerializedBaseOperator._serialize(real_op) @@ -1590,14 +1593,58 @@ def test_mapped_operator_serde(): {"__type": "dict", "__var": {'a': 'b'}}, ] }, - 'partial_kwargs': {}, + 'partial_kwargs': { + 'executor_config': { + '__type': 'dict', + '__var': { + 'dict': {"__type": "dict", "__var": {'sub': 'value'}}, + }, + }, + }, 'task_id': 'a', 'template_fields': ['bash_command', 'env'], } op = SerializedBaseOperator.deserialize_operator(serialized) + assert isinstance(op, MappedOperator) assert op.operator_class == "airflow.operators.bash.BashOperator" + assert op.mapped_kwargs['bash_command'] == literal + assert op.partial_kwargs['executor_config'] == {'dict': {'sub': 'value'}} + + +def test_mapped_operator_xcomarg_serde(): + from airflow.models.xcom_arg import XComArg + + with DAG("test-dag", start_date=datetime(2020, 1, 1)) as dag: + task1 = BaseOperator(task_id="op1") + xcomarg = XComArg(task1, "test_key") + mapped = MockOperator(task_id='task_2').map(arg2=xcomarg) + + serialized = SerializedBaseOperator._serialize(mapped) + assert serialized == { + '_is_dummy': False, + '_is_mapped': True, + '_task_module': 'tests.test_utils.mock_operators', + '_task_type': 'MockOperator', + 'downstream_task_ids': [], + 'mapped_kwargs': {'arg2': {'__type': 'xcomref', '__var': {'task_id': 'op1', 'key': 'test_key'}}}, + 'partial_kwargs': {}, + 'task_id': 'task_2', + 'template_fields': ['arg1', 'arg2'], + } + + op = SerializedBaseOperator.deserialize_operator(serialized) + + arg = op.mapped_kwargs['arg2'] + assert arg.task_id == 'op1' + assert arg.key == 'test_key' + + serialized_dag: DAG = SerializedDAG.from_dict(SerializedDAG.to_dict(dag)) + + xcom_arg = serialized_dag.task_dict['task_2'].mapped_kwargs['arg2'] + assert isinstance(xcom_arg, XComArg) + assert xcom_arg.operator is serialized_dag.task_dict['op1'] def test_mapped_task_group_serde():