diff --git a/airflow/decorators/base.py b/airflow/decorators/base.py index e8216625758af..2eafaaab54144 100644 --- a/airflow/decorators/base.py +++ b/airflow/decorators/base.py @@ -277,24 +277,18 @@ def _validate_arg_names(self, funcname: str, kwargs: Dict[str, Any], valid_names def map( self, *, dag: Optional["DAG"] = None, task_group: Optional["TaskGroup"] = None, **kwargs ) -> XComArg: - + self._validate_arg_names("map", kwargs) dag = dag or DagContext.get_current_dag() task_group = task_group or TaskGroupContext.get_current_task_group(dag) task_id = get_unique_task_id(self.kwargs['task_id'], dag, task_group) - self._validate_arg_names("map", kwargs) - - operator = MappedOperator( - operator_class=self.operator_class, - task_id=task_id, + operator = MappedOperator.from_decorator( + decorator=self, dag=dag, task_group=task_group, - partial_kwargs=self.kwargs, - # Set them to empty to bypass the validation, as for decorated stuff we validate ourselves - mapped_kwargs={}, + task_id=task_id, + mapped_kwargs=kwargs, ) - operator.mapped_kwargs.update(kwargs) - return XComArg(operator=operator) def partial( diff --git a/airflow/exceptions.py b/airflow/exceptions.py index d1772153c6ab6..96ad39f4adc55 100644 --- a/airflow/exceptions.py +++ b/airflow/exceptions.py @@ -99,6 +99,17 @@ class AirflowFailException(AirflowException): """Raise when the task should be failed without retrying.""" +class UnmappableXComPushed(AirflowException): + """Raise when an unmappable value is pushed as a mapped downstream's dependency.""" + + def __init__(self, value: Any) -> None: + super().__init__(value) + self.value = value + + def __str__(self) -> str: + return f"unmappable return type {type(self.value).__qualname__!r}" + + class AirflowDagCycleException(AirflowException): """Raise when there is a cycle in DAG definition.""" diff --git a/airflow/migrations/versions/e655c0453f75_add_taskmap_and_map_id_on_taskinstance.py b/airflow/migrations/versions/e655c0453f75_add_taskmap_and_map_id_on_taskinstance.py new file mode 100644 index 0000000000000..91afcb99e6f39 --- /dev/null +++ b/airflow/migrations/versions/e655c0453f75_add_taskmap_and_map_id_on_taskinstance.py @@ -0,0 +1,120 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Add TaskMap and map_index on TaskInstance. + +Revision ID: e655c0453f75 +Revises: 587bdf053233 +Create Date: 2021-12-13 22:59:41.052584 +""" + +from alembic import op +from sqlalchemy import Column, ForeignKeyConstraint, Integer + +from airflow.models.base import StringID +from airflow.utils.sqlalchemy import ExtendedJSON + +# Revision identifiers, used by Alembic. +revision = "e655c0453f75" +down_revision = "587bdf053233" +branch_labels = None +depends_on = None + + +def upgrade(): + """Add TaskMap and map_index on TaskInstance.""" + # We need to first remove constraints on task_reschedule since they depend on task_instance. + with op.batch_alter_table("task_reschedule") as batch_op: + batch_op.drop_constraint("task_reschedule_ti_fkey", "foreignkey") + batch_op.drop_index("idx_task_reschedule_dag_task_run") + + # Change task_instance's primary key. + with op.batch_alter_table("task_instance") as batch_op: + # I think we always use this name for TaskInstance after 7b2661a43ba3? + batch_op.drop_constraint("task_instance_pkey", type_="primary") + batch_op.add_column(Column("map_index", Integer, nullable=False, default=-1)) + batch_op.create_primary_key("task_instance_pkey", ["dag_id", "task_id", "run_id", "map_index"]) + + # Re-create task_reschedule's constraints. + with op.batch_alter_table("task_reschedule") as batch_op: + batch_op.add_column(Column("map_index", Integer, nullable=False, default=-1)) + batch_op.create_foreign_key( + "task_reschedule_ti_fkey", + "task_instance", + ["dag_id", "task_id", "run_id", "map_index"], + ["dag_id", "task_id", "run_id", "map_index"], + ondelete="CASCADE", + ) + batch_op.create_index( + "idx_task_reschedule_dag_task_run", + ["dag_id", "task_id", "run_id", "map_index"], + unique=False, + ) + + # Create task_map. + op.create_table( + "task_map", + Column("dag_id", StringID(), primary_key=True), + Column("task_id", StringID(), primary_key=True), + Column("run_id", StringID(), primary_key=True), + Column("map_index", Integer, primary_key=True), + Column("length", Integer, nullable=False), + Column("keys", ExtendedJSON, nullable=True), + ForeignKeyConstraint( + ["dag_id", "task_id", "run_id", "map_index"], + [ + "task_instance.dag_id", + "task_instance.task_id", + "task_instance.run_id", + "task_instance.map_index", + ], + name="task_map_task_instance_fkey", + ondelete="CASCADE", + ), + ) + + +def downgrade(): + """Remove TaskMap and map_index on TaskInstance.""" + op.drop_table("task_map") + + with op.batch_alter_table("task_reschedule") as batch_op: + batch_op.drop_constraint("task_reschedule_ti_fkey", "foreignkey") + batch_op.drop_index("idx_task_reschedule_dag_task_run") + batch_op.drop_column("map_index") + + op.execute("DELETE FROM task_instance WHERE map_index != -1") + + with op.batch_alter_table("task_instance") as batch_op: + batch_op.drop_constraint("task_instance_pkey", type_="primary") + batch_op.drop_column("map_index") + batch_op.create_primary_key("task_instance_pkey", ["dag_id", "task_id", "run_id"]) + + with op.batch_alter_table("task_reschedule") as batch_op: + batch_op.create_foreign_key( + "task_reschedule_ti_fkey", + "task_instance", + ["dag_id", "task_id", "run_id"], + ["dag_id", "task_id", "run_id"], + ondelete="CASCADE", + ) + batch_op.create_index( + "idx_task_reschedule_dag_task_run", + ["dag_id", "task_id", "run_id"], + unique=False, + ) diff --git a/airflow/models/baseoperator.py b/airflow/models/baseoperator.py index e1eee43317f7d..3b9d792223210 100644 --- a/airflow/models/baseoperator.py +++ b/airflow/models/baseoperator.py @@ -78,6 +78,7 @@ from airflow.utils.weight_rule import WeightRule if TYPE_CHECKING: + from airflow.decorators.base import _TaskDecorator from airflow.models.dag import DAG from airflow.utils.task_group import TaskGroup @@ -1632,6 +1633,39 @@ def defer( def map(self, **kwargs) -> "MappedOperator": return MappedOperator.from_operator(self, kwargs) + def has_mapped_dependants(self) -> bool: + """Whether any downstream dependencies depend on this task for mapping. + + For now, this walks the entire DAG to find mapped nodes that has this + current task as an upstream. We cannot use ``downstream_list`` since it + only contains operators, not task groups. In the future, we should + provide a way to record an DAG node's all downstream nodes instead. + """ + from airflow.utils.task_group import MappedTaskGroup, TaskGroup + + if not self.has_dag(): + return False + + def _walk_group(group: TaskGroup) -> Iterable[Tuple[str, DAGNode]]: + """Recursively walk children in a task group. + + This yields all direct children (including both tasks and task + groups), and all children of any task groups. + """ + for key, child in group.children.items(): + yield key, child + if isinstance(child, TaskGroup): + yield from _walk_group(child) + + for key, child in _walk_group(self.dag.task_group): + if key == self.task_id: + continue + if not isinstance(child, (MappedOperator, MappedTaskGroup)): + continue + if self.task_id in child.upstream_task_ids: + return True + return False + def _validate_kwarg_names_for_mapping(cls: Type[BaseOperator], func_name: str, value: Dict[str, Any]): if isinstance(str, cls): @@ -1668,6 +1702,7 @@ class MappedOperator(DAGNode): """Object representing a mapped operator in a DAG""" operator_class: Type[BaseOperator] = attr.ib(repr=lambda c: c.__name__) + task_type: str = attr.ib() task_id: str partial_kwargs: Dict[str, Any] mapped_kwargs: Dict[str, Any] = attr.ib( @@ -1678,9 +1713,12 @@ class MappedOperator(DAGNode): downstream_task_ids: Set[str] = attr.ib(factory=set, repr=False) task_group: Optional["TaskGroup"] = attr.ib(repr=False) + # BaseOperator-like interface -- needed so we can add oursleves to the dag.tasks start_date: Optional[pendulum.DateTime] = attr.ib(repr=False, default=None) end_date: Optional[pendulum.DateTime] = attr.ib(repr=False, default=None) + owner: str = attr.ib(repr=False, default=conf.get("operators", "DEFAULT_OWNER")) + max_active_tis_per_dag: Optional[int] = attr.ib(default=None) @classmethod def from_operator(cls, operator: BaseOperator, mapped_kwargs: Dict[str, Any]) -> "MappedOperator": @@ -1699,8 +1737,36 @@ def from_operator(cls, operator: BaseOperator, mapped_kwargs: Dict[str, Any]) -> end_date=operator.end_date, partial_kwargs=operator._BaseOperator__init_kwargs, # type: ignore mapped_kwargs=mapped_kwargs, + owner=operator.owner, + max_active_tis_per_dag=operator.max_active_tis_per_dag, ) + @classmethod + def from_decorator( + cls, + *, + decorator: "_TaskDecorator", + dag: Optional["DAG"], + task_group: Optional["TaskGroup"], + task_id: str, + mapped_kwargs: Dict[str, Any], + ) -> "MappedOperator": + """Create a mapped operator from a task decorator. + + Different from ``from_operator``, this DOES NOT validate ``mapped_kwargs``. + The task decorator calling this should be responsible for validation. + """ + operator = MappedOperator( + operator_class=decorator.operator_class, + partial_kwargs=decorator.kwargs, + mapped_kwargs={}, + task_id=task_id, + dag=dag, + task_group=task_group, + ) + operator.mapped_kwargs.update(mapped_kwargs) + return operator + def __attrs_post_init__(self): if self.task_group: self.task_id = self.task_group.child_id(self.task_id) @@ -1708,6 +1774,10 @@ def __attrs_post_init__(self): if self.dag: self.dag.add_task(self) + @task_type.default + def _default_task_type(self): + return self.operator_class.__name__ + @task_group.default def _default_task_group(self): from airflow.utils.task_group import TaskGroupContext diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index 5015f4626b1f8..0382ad5ef8792 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -15,6 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +import collections.abc import contextlib import hashlib import logging @@ -82,11 +83,13 @@ DagRunNotFound, TaskDeferralError, TaskDeferred, + UnmappableXComPushed, ) from airflow.models.base import COLLATION_ARGS, ID_LEN, Base from airflow.models.log import Log from airflow.models.param import ParamsDict from airflow.models.taskfail import TaskFail +from airflow.models.taskmap import TaskMap from airflow.models.taskreschedule import TaskReschedule from airflow.models.xcom import XCOM_RETURN_KEY, XCom from airflow.plugins_manager import integrate_macros_plugins @@ -340,6 +343,8 @@ class TaskInstance(Base, LoggingMixin): task_id = Column(String(ID_LEN, **COLLATION_ARGS), primary_key=True, nullable=False) dag_id = Column(String(ID_LEN, **COLLATION_ARGS), primary_key=True, nullable=False) run_id = Column(String(ID_LEN, **COLLATION_ARGS), primary_key=True, nullable=False) + map_index = Column(Integer, primary_key=True, nullable=False, default=-1) + start_date = Column(UtcDateTime) end_date = Column(UtcDateTime) duration = Column(Float) @@ -424,10 +429,12 @@ def __init__( execution_date: Optional[datetime] = None, run_id: Optional[str] = None, state: Optional[str] = None, + map_index: int = -1, ): super().__init__() self.dag_id = task.dag_id self.task_id = task.task_id + self.map_index = map_index self.refresh_from_task(task) self._log = logging.getLogger("airflow.task") @@ -760,6 +767,7 @@ def refresh_from_db(self, session=NEW_SESSION, lock_for_update=False) -> None: TaskInstance.dag_id == self.dag_id, TaskInstance.task_id == self.task_id, TaskInstance.run_id == self.run_id, + TaskInstance.map_index == self.map_index, ) if lock_for_update: @@ -1537,7 +1545,9 @@ def _execute_task(self, context, task_copy): result = execute_callable(context=context) # If the task returns a result, push an XCom containing it if task_copy.do_xcom_push and result is not None: - self.xcom_push(key=XCOM_RETURN_KEY, value=result) + with create_session() as session: + self.xcom_push(key=XCOM_RETURN_KEY, value=result, session=session) + self._record_task_map_for_downstreams(result, session=session) return result @provide_session @@ -2128,6 +2138,14 @@ def set_duration(self) -> None: self.duration = None self.log.debug("Task Duration set to %s", self.duration) + def _record_task_map_for_downstreams(self, value: Any, *, session: Session) -> None: + if not self.task.has_mapped_dependants(): + return + if not isinstance(value, collections.abc.Collection) or isinstance(value, (bytes, str)): + self.log.info("Failing %s for unmappable XCom push %r", self.key, type(value).__qualname__) + raise UnmappableXComPushed(value) + session.merge(TaskMap.from_task_instance_xcom(self, value)) + @provide_session def xcom_push( self, diff --git a/airflow/models/taskmap.py b/airflow/models/taskmap.py new file mode 100644 index 0000000000000..e91fbd6148c29 --- /dev/null +++ b/airflow/models/taskmap.py @@ -0,0 +1,110 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Table to store information about mapped task instances (AIP-42).""" + +import collections.abc +import enum +from typing import TYPE_CHECKING, Any, Collection, List, Optional + +from sqlalchemy import Column, ForeignKeyConstraint, Integer, String + +from airflow.models.base import COLLATION_ARGS, ID_LEN, Base +from airflow.utils.sqlalchemy import ExtendedJSON + +if TYPE_CHECKING: + from airflow.models.taskinstance import TaskInstance + + +class TaskMapVariant(enum.Enum): + """Task map variant. + + Possible values are **dict** (for a key-value mapping) and **list** (for an + ordered value sequence). + """ + + DICT = "dict" + LIST = "list" + + +class TaskMap(Base): + """Model to track dynamic task-mapping information. + + This is currently only populated by an upstream TaskInstance pushing an + XCom that's pulled by a downstream for mapping purposes. + """ + + __tablename__ = "task_map" + + # Link to upstream TaskInstance creating this dynamic mapping information. + dag_id = Column(String(ID_LEN, **COLLATION_ARGS), primary_key=True) + task_id = Column(String(ID_LEN, **COLLATION_ARGS), primary_key=True) + run_id = Column(String(ID_LEN, **COLLATION_ARGS), primary_key=True) + map_index = Column(Integer, primary_key=True) + + length = Column(Integer, nullable=False) + keys = Column(ExtendedJSON, nullable=True) + + __table_args__ = ( + ForeignKeyConstraint( + [dag_id, task_id, run_id, map_index], + [ + "task_instance.dag_id", + "task_instance.task_id", + "task_instance.run_id", + "task_instance.map_index", + ], + name="task_map_task_instance_fkey", + ondelete="CASCADE", + ), + ) + + def __init__( + self, + dag_id: str, + task_id: str, + run_id: str, + map_index: int, + length: int, + keys: Optional[List[Any]], + ) -> None: + self.dag_id = dag_id + self.task_id = task_id + self.run_id = run_id + self.map_index = map_index + self.length = length + self.keys = keys + + @classmethod + def from_task_instance_xcom(cls, ti: "TaskInstance", value: Collection) -> "TaskMap": + if ti.run_id is None: + raise ValueError("cannot record task map for unrun task instance") + return cls( + dag_id=ti.dag_id, + task_id=ti.task_id, + run_id=ti.run_id, + map_index=ti.map_index, + length=len(value), + keys=(list(value) if isinstance(value, collections.abc.Mapping) else None), + ) + + @property + def variant(self) -> TaskMapVariant: + if self.keys is None: + return TaskMapVariant.LIST + return TaskMapVariant.DICT diff --git a/airflow/models/taskreschedule.py b/airflow/models/taskreschedule.py index 55ef754da27ff..4e39af5beea4f 100644 --- a/airflow/models/taskreschedule.py +++ b/airflow/models/taskreschedule.py @@ -16,6 +16,10 @@ # specific language governing permissions and limitations # under the License. """TaskReschedule tracks rescheduled task instances.""" + +import datetime +from typing import TYPE_CHECKING + from sqlalchemy import Column, ForeignKeyConstraint, Index, Integer, String, asc, desc from sqlalchemy.ext.associationproxy import association_proxy from sqlalchemy.orm import relationship @@ -24,6 +28,9 @@ from airflow.utils.session import provide_session from airflow.utils.sqlalchemy import UtcDateTime +if TYPE_CHECKING: + from airflow.models.baseoperator import BaseOperator + class TaskReschedule(Base): """TaskReschedule tracks rescheduled task instances.""" @@ -34,6 +41,7 @@ class TaskReschedule(Base): task_id = Column(String(ID_LEN, **COLLATION_ARGS), nullable=False) dag_id = Column(String(ID_LEN, **COLLATION_ARGS), nullable=False) run_id = Column(String(ID_LEN, **COLLATION_ARGS), nullable=False) + map_index = Column(Integer, nullable=False, default=-1) try_number = Column(Integer, nullable=False) start_date = Column(UtcDateTime, nullable=False) end_date = Column(UtcDateTime, nullable=False) @@ -41,12 +49,17 @@ class TaskReschedule(Base): reschedule_date = Column(UtcDateTime, nullable=False) __table_args__ = ( - Index('idx_task_reschedule_dag_task_run', dag_id, task_id, run_id, unique=False), + Index('idx_task_reschedule_dag_task_run', dag_id, task_id, run_id, map_index, unique=False), ForeignKeyConstraint( - [dag_id, task_id, run_id], - ['task_instance.dag_id', 'task_instance.task_id', 'task_instance.run_id'], - name='task_reschedule_ti_fkey', - ondelete='CASCADE', + [dag_id, task_id, run_id, map_index], + [ + "task_instance.dag_id", + "task_instance.task_id", + "task_instance.run_id", + "task_instance.map_index", + ], + name="task_reschedule_ti_fkey", + ondelete="CASCADE", ), ForeignKeyConstraint( [dag_id, run_id], @@ -58,10 +71,20 @@ class TaskReschedule(Base): dag_run = relationship("DagRun") execution_date = association_proxy("dag_run", "execution_date") - def __init__(self, task, run_id, try_number, start_date, end_date, reschedule_date): + def __init__( + self, + task: "BaseOperator", + run_id: str, + try_number: int, + start_date: datetime.datetime, + end_date: datetime.datetime, + reschedule_date: datetime.datetime, + map_index: int = -1, + ): self.dag_id = task.dag_id self.task_id = task.task_id self.run_id = run_id + self.map_index = map_index self.try_number = try_number self.start_date = start_date self.end_date = end_date diff --git a/docs/apache-airflow/migrations-ref.rst b/docs/apache-airflow/migrations-ref.rst index 5c9e24175d943..f004b343fa522 100644 --- a/docs/apache-airflow/migrations-ref.rst +++ b/docs/apache-airflow/migrations-ref.rst @@ -23,7 +23,10 @@ Here's the list of all the Database Migrations that are executed via when you ru +--------------------------------+------------------+-----------------+---------------------------------------------------------------------------------------+ | Revision ID | Revises ID | Airflow Version | Description | +--------------------------------+------------------+-----------------+---------------------------------------------------------------------------------------+ -| ``587bdf053233`` (head) | ``f9da662e7089`` | ``2.3.0`` | Add index for ``dag_id`` column in ``job`` table. | +| ``e655c0453f75`` (head) | ``587bdf053233`` | ``2.3.0`` | Add ``map_index`` column to TaskInstance to identify task-mapping, and a ``task_map`` | +| | | | table to track mapping values from XCom. | ++--------------------------------+------------------+-----------------+---------------------------------------------------------------------------------------+ +| ``587bdf053233`` | ``f9da662e7089`` | ``2.3.0`` | Add index for ``dag_id`` column in ``job`` table. | +--------------------------------+------------------+-----------------+---------------------------------------------------------------------------------------+ | ``f9da662e7089`` | ``786e3737b18f`` | ``2.3.0`` | Add ``LogTemplate`` table to track changes to config values ``log_filename_template`` | +--------------------------------+------------------+-----------------+---------------------------------------------------------------------------------------+ diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index eebe41da31275..9ebf07d41dda7 100644 --- a/docs/spelling_wordlist.txt +++ b/docs/spelling_wordlist.txt @@ -1408,6 +1408,7 @@ unicode unittest unittests unix +unmappable unmapped unpause unpausing diff --git a/tests/conftest.py b/tests/conftest.py index bfea15d0ab56d..3f68c9e318e99 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -597,6 +597,7 @@ def __call__( def cleanup(self): from airflow.models import DagModel, DagRun, TaskInstance, XCom from airflow.models.serialized_dag import SerializedDagModel + from airflow.models.taskmap import TaskMap from airflow.utils.retries import run_with_db_retries for attempt in run_with_db_retries(logger=self.log): @@ -611,16 +612,19 @@ def cleanup(self): SerializedDagModel.dag_id.in_(dag_ids) ).delete(synchronize_session=False) self.session.query(DagRun).filter(DagRun.dag_id.in_(dag_ids)).delete( - synchronize_session=False + synchronize_session=False, ) self.session.query(TaskInstance).filter(TaskInstance.dag_id.in_(dag_ids)).delete( - synchronize_session=False + synchronize_session=False, ) self.session.query(XCom).filter(XCom.dag_id.in_(dag_ids)).delete( - synchronize_session=False + synchronize_session=False, ) self.session.query(DagModel).filter(DagModel.dag_id.in_(dag_ids)).delete( - synchronize_session=False + synchronize_session=False, + ) + self.session.query(TaskMap).filter(TaskMap.dag_id.in_(dag_ids)).delete( + synchronize_session=False, ) self.session.commit() if self._own_session: diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py index 4b6a48cd4a21b..e11594f75fcc0 100644 --- a/tests/models/test_taskinstance.py +++ b/tests/models/test_taskinstance.py @@ -37,6 +37,7 @@ AirflowRescheduleException, AirflowSensorTimeout, AirflowSkipException, + UnmappableXComPushed, ) from airflow.models import ( DAG, @@ -50,6 +51,7 @@ XCom, ) from airflow.models.taskinstance import load_error_file, set_error_file +from airflow.models.taskmap import TaskMap from airflow.operators.bash import BashOperator from airflow.operators.dummy import DummyOperator from airflow.operators.python import PythonOperator @@ -2060,6 +2062,7 @@ def test_refresh_from_db(self, create_task_instance): "task_id": "test_refresh_from_db_task", "dag_id": "test_refresh_from_db_dag", "run_id": "test", + "map_index": -1, "start_date": run_date + datetime.timedelta(days=1), "end_date": run_date + datetime.timedelta(days=1, seconds=1, milliseconds=234), "duration": 1.234, @@ -2085,7 +2088,7 @@ def test_refresh_from_db(self, create_task_instance): "next_method": None, } # Make sure we aren't missing any new value in our expected_values list. - expected_keys = {f"task_instance.{key.lstrip('_')}" for key in expected_values.keys()} + expected_keys = {f"task_instance.{key.lstrip('_')}" for key in expected_values} assert {str(c) for c in TI.__table__.columns} == expected_keys, ( "Please add all non-foreign values of TaskInstance to this list. " "This prevents refresh_from_db() from missing a field." @@ -2224,7 +2227,7 @@ def timeout(): mock_on_failure = mock.MagicMock() with dag_maker(dag_id=f'test_sensor_timeout_{mode}_{retries}'): - task = PythonSensor( + PythonSensor( task_id='test_raise_sensor_timeout', python_callable=timeout, on_failure_callback=mock_on_failure, @@ -2232,10 +2235,88 @@ def timeout(): mode=mode, ) ti = dag_maker.create_dagrun(execution_date=timezone.utcnow()).task_instances[0] - ti.task = task with pytest.raises(AirflowSensorTimeout): ti.run() assert mock_on_failure.called assert ti.state == State.FAILED + + +class TestTaskInstanceRecordTaskMapXComPush: + """Test TI.xcom_push() correctly records return values for task-mapping.""" + + def setup_class(self): + """Ensure we start fresh.""" + with create_session() as session: + session.query(TaskMap).delete() + + def _run_ti_with_faked_mapped_dependants(self, ti): + # TODO: We can't actually put a MappedOperator in a DAG yet due to it + # lacking some functions we expect from BaseOperator, so we mock this + # instead to test what effect it has to TaskMap recording. + with mock.patch.object(ti.task, "has_mapped_dependants", new=lambda: True): + ti.run() + + @pytest.mark.parametrize("xcom_value", [[1, 2, 3], {"a": 1, "b": 2}, "abc"]) + def test_not_recorded_for_unused(self, dag_maker, xcom_value): + """A value not used for task-mapping should not be recorded.""" + with dag_maker(dag_id="test_not_recorded_for_unused") as dag: + + @dag.task() + def push_something(): + return xcom_value + + push_something() + + ti = dag_maker.create_dagrun().task_instances[0] + ti.run() + + assert dag_maker.session.query(TaskMap).count() == 0 + + def test_error_if_unmappable(self, caplog, dag_maker): + """If an unmappable return value is used to map, fail the task that pushed the XCom.""" + with dag_maker(dag_id="test_not_recorded_for_unused") as dag: + + @dag.task() + def push_something(): + return "abc" + + push_something() + + ti = dag_maker.create_dagrun().task_instances[0] + with pytest.raises(UnmappableXComPushed) as ctx: + self._run_ti_with_faked_mapped_dependants(ti) + + assert dag_maker.session.query(TaskMap).count() == 0 + assert ti.state == TaskInstanceState.FAILED + assert str(ctx.value) == "unmappable return type 'str'" + + @pytest.mark.parametrize( + "xcom_value, expected_length, expected_keys", + [ + ([1, 2, 3], 3, None), + ({"a": 1, "b": 2}, 2, ["a", "b"]), + ], + ) + def test_written_task_map(self, dag_maker, xcom_value, expected_length, expected_keys): + """Return value should be recorded in TaskMap if it's used by a downstream to map.""" + with dag_maker(dag_id="test_written_task_map") as dag: + + @dag.task() + def push_something(): + return xcom_value + + push_something() + + dag_run = dag_maker.create_dagrun() + ti = next(ti for ti in dag_run.task_instances if ti.task_id == "push_something") + self._run_ti_with_faked_mapped_dependants(ti) + + task_map = dag_maker.session.query(TaskMap).one() + assert task_map.dag_id == "test_written_task_map" + assert task_map.task_id == "push_something" + assert task_map.run_id == dag_run.run_id + assert task_map.map_index == -1 + assert task_map.length == expected_length + assert task_map.keys == expected_keys