From c8ba1ec0f057fc135e26cec02a4591a36564dead Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Tue, 14 Dec 2021 09:30:28 +0800 Subject: [PATCH 01/20] Add TaskMap and TaskInstance.map_id TaskMap is populated by a TaskInstance when it pushes its return value to XCom. --- ..._add_taskmap_and_map_id_on_taskinstance.py | 76 +++++++++++++++++++ airflow/models/taskinstance.py | 31 ++++++++ airflow/models/taskmap.py | 76 +++++++++++++++++++ docs/apache-airflow/migrations-ref.rst | 5 +- 4 files changed, 187 insertions(+), 1 deletion(-) create mode 100644 airflow/migrations/versions/e655c0453f75_add_taskmap_and_map_id_on_taskinstance.py create mode 100644 airflow/models/taskmap.py 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..1af685e343e81 --- /dev/null +++ b/airflow/migrations/versions/e655c0453f75_add_taskmap_and_map_id_on_taskinstance.py @@ -0,0 +1,76 @@ +# +# 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, String, text + +from airflow.models.base import COLLATION_ARGS +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.""" + with op.batch_alter_table("task_instance") as batch_op: + batch_op.add_column(Column("map_index", Integer, server_default=text("-1"))) + # I think we always use this name for TaskInstance after 7b2661a43ba3? + batch_op.drop_constraint("task_instance_pkey", type_="primary") + batch_op.create_primary_key("task_instance_pkey", ["dag_id", "task_id", "run_id", "map_index"]) + + op.create_table( + "task_map", + Column("dag_id", String(250, **COLLATION_ARGS), primary_key=True), + Column("task_id", String(250, **COLLATION_ARGS), primary_key=True), + Column("run_id", String(250, **COLLATION_ARGS), 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") + 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.create_primary_key("task_instance_pkey", ["dag_id", "task_id", "run_id"]) + batch_op.drop_column("map_index") diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index 5015f4626b1f8..540461b6805fe 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 @@ -59,6 +60,7 @@ func, inspect, or_, + text, tuple_, ) from sqlalchemy.ext.associationproxy import association_proxy @@ -87,6 +89,7 @@ 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, server_default=text("-1")) + start_date = Column(UtcDateTime) end_date = Column(UtcDateTime) duration = Column(Float) @@ -2128,6 +2133,29 @@ def set_duration(self) -> None: self.duration = None self.log.debug("Task Duration set to %s", self.duration) + @provide_session + def _record_task_map_for_downstreams(self, value: Any, *, session: Session = NEW_SESSION) -> None: + # TODO: Only record if we know this value is used in that mapped + # downstream task or task group. We should also error if there are + # mapped downstreams, but the pushed value is not mappable. + if not isinstance(value, collections.abc.Collection): + return + session.query(TaskMap).filter_by( + dag_id=self.dag_id, + task_id=self.task_id, + run_id=self.run_id, + map_index=self.map_index, + ).delete() + instance = TaskMap( + dag_id=self.dag_id, + task_id=self.task_id, + run_id=self.run_id, + map_index=self.map_index, + length=len(value), + keys=(list(value) if isinstance(value, collections.abc.Mapping) else None), + ) + session.merge(instance) + @provide_session def xcom_push( self, @@ -2167,6 +2195,9 @@ def xcom_push( session=session, ) + if key == XCOM_RETURN_KEY: + self._record_task_map_for_downstreams(value, session=session) + @provide_session def xcom_pull( self, diff --git a/airflow/models/taskmap.py b/airflow/models/taskmap.py new file mode 100644 index 0000000000000..6ff95a35e14a8 --- /dev/null +++ b/airflow/models/taskmap.py @@ -0,0 +1,76 @@ +# +# 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 "mapped" task instances (AIP-42).""" + +import enum + +from sqlalchemy import Column, ForeignKeyConstraint, Integer, String + +from airflow.models.base import COLLATION_ARGS, ID_LEN, Base +from airflow.utils.sqlalchemy import ExtendedJSON + + +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__ = "xcom_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", + ), + ) + + @property + def variant(self) -> TaskMapVariant: + if self.keys is None: + return TaskMapVariant.LIST + return TaskMapVariant.DICT 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`` | +--------------------------------+------------------+-----------------+---------------------------------------------------------------------------------------+ From 8a66a1bf558f3404de64e2a5ed889b33c7b806cf Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Tue, 14 Dec 2021 11:41:54 +0800 Subject: [PATCH 02/20] Only push task map when needed as dependencies --- airflow/models/baseoperator.py | 27 +++++++++++++++++++++++++++ airflow/models/taskinstance.py | 7 +++---- 2 files changed, 30 insertions(+), 4 deletions(-) diff --git a/airflow/models/baseoperator.py b/airflow/models/baseoperator.py index e1eee43317f7d..1114eef486641 100644 --- a/airflow/models/baseoperator.py +++ b/airflow/models/baseoperator.py @@ -1632,6 +1632,33 @@ def defer( def map(self, **kwargs) -> "MappedOperator": return MappedOperator.from_operator(self, kwargs) + def has_mapped_dependants(self) -> bool: + """Whether any downstreams depend on this task for mapping.""" + from airflow.utils.task_group import MappedTaskGroup + + 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): diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index 540461b6805fe..72ed418bd0650 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -2135,11 +2135,10 @@ def set_duration(self) -> None: @provide_session def _record_task_map_for_downstreams(self, value: Any, *, session: Session = NEW_SESSION) -> None: - # TODO: Only record if we know this value is used in that mapped - # downstream task or task group. We should also error if there are - # mapped downstreams, but the pushed value is not mappable. - if not isinstance(value, collections.abc.Collection): + if not self.task.has_mapped_dependants(): return + if not isinstance(value, collections.abc.Collection): + return # TODO: Error if the pushed value is not mappable? session.query(TaskMap).filter_by( dag_id=self.dag_id, task_id=self.task_id, From d7e7450225663d1aff47f90cd34e9204a946bd6e Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Tue, 14 Dec 2021 12:05:31 +0800 Subject: [PATCH 03/20] Fix task_reschedule migration --- ..._add_taskmap_and_map_id_on_taskinstance.py | 52 ++++++++++++++++++- airflow/models/taskreschedule.py | 18 ++++--- 2 files changed, 62 insertions(+), 8 deletions(-) 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 index 1af685e343e81..c8f969e84d6d5 100644 --- 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 @@ -38,12 +38,37 @@ 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_index("idx_task_reschedule_dag_task_run") + batch_op.drop_constraint("task_reschedule_ti_fkey", "foreignkey") + if op.get_bind().dialect.name == "mysql": # MySQL also creates an index. + batch_op.drop_index("task_reschedule_ti_fkey") + + # Change task_instance's primary key. with op.batch_alter_table("task_instance") as batch_op: - batch_op.add_column(Column("map_index", Integer, server_default=text("-1"))) # 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, server_default=text("-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, server_default=text("-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", String(250, **COLLATION_ARGS), primary_key=True), @@ -69,8 +94,31 @@ def upgrade(): 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_index("idx_task_reschedule_dag_task_run") + batch_op.drop_constraint("task_reschedule_ti_fkey", "foreignkey") + if op.get_bind().dialect.name == "mysql": # MySQL also creates an index. + batch_op.drop_index("task_reschedule_ti_fkey") + 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.create_primary_key("task_instance_pkey", ["dag_id", "task_id", "run_id"]) 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/taskreschedule.py b/airflow/models/taskreschedule.py index 55ef754da27ff..246b0bcd5953c 100644 --- a/airflow/models/taskreschedule.py +++ b/airflow/models/taskreschedule.py @@ -16,7 +16,7 @@ # specific language governing permissions and limitations # under the License. """TaskReschedule tracks rescheduled task instances.""" -from sqlalchemy import Column, ForeignKeyConstraint, Index, Integer, String, asc, desc +from sqlalchemy import Column, ForeignKeyConstraint, Index, Integer, String, asc, desc, text from sqlalchemy.ext.associationproxy import association_proxy from sqlalchemy.orm import relationship @@ -34,6 +34,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, server_default=text("-1")) try_number = Column(Integer, nullable=False) start_date = Column(UtcDateTime, nullable=False) end_date = Column(UtcDateTime, nullable=False) @@ -41,12 +42,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], From d97c086dd2411e4bb87e11527a8000cdc5703cd6 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Tue, 14 Dec 2021 14:18:59 +0800 Subject: [PATCH 04/20] Fix TaskMap table name --- airflow/models/taskmap.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/models/taskmap.py b/airflow/models/taskmap.py index 6ff95a35e14a8..17f7b6f4f589b 100644 --- a/airflow/models/taskmap.py +++ b/airflow/models/taskmap.py @@ -44,7 +44,7 @@ class TaskMap(Base): XCom that's pulled by a downstream for mapping purposes. """ - __tablename__ = "xcom_task_map" + __tablename__ = "task_map" # Link to upstream TaskInstance creating this dynamic mapping information. dag_id = Column(String(ID_LEN, **COLLATION_ARGS), primary_key=True) From 74dd354099b17d2c60954b142664f0389bcce46a Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Tue, 14 Dec 2021 23:43:52 +0800 Subject: [PATCH 05/20] Must re-import TaskGrounp at runtime --- airflow/models/baseoperator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/models/baseoperator.py b/airflow/models/baseoperator.py index 1114eef486641..7fd5a5136c5fd 100644 --- a/airflow/models/baseoperator.py +++ b/airflow/models/baseoperator.py @@ -1634,7 +1634,7 @@ def map(self, **kwargs) -> "MappedOperator": def has_mapped_dependants(self) -> bool: """Whether any downstreams depend on this task for mapping.""" - from airflow.utils.task_group import MappedTaskGroup + from airflow.utils.task_group import MappedTaskGroup, TaskGroup if not self.has_dag(): return False From fff3a1e866bfc2d00d6454a5feb72203a8ce7513 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Wed, 15 Dec 2021 03:21:10 +0800 Subject: [PATCH 06/20] Use client-side default --- .../e655c0453f75_add_taskmap_and_map_id_on_taskinstance.py | 6 +++--- airflow/models/taskinstance.py | 3 +-- airflow/models/taskreschedule.py | 4 ++-- 3 files changed, 6 insertions(+), 7 deletions(-) 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 index c8f969e84d6d5..119f926d3614a 100644 --- 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 @@ -24,7 +24,7 @@ """ from alembic import op -from sqlalchemy import Column, ForeignKeyConstraint, Integer, String, text +from sqlalchemy import Column, ForeignKeyConstraint, Integer, String from airflow.models.base import COLLATION_ARGS from airflow.utils.sqlalchemy import ExtendedJSON @@ -49,12 +49,12 @@ def upgrade(): 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, server_default=text("-1"))) + batch_op.add_column(Column("map_index", Integer, 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, server_default=text("-1"))) + batch_op.add_column(Column("map_index", Integer, nullable=False, default=-1)) batch_op.create_foreign_key( "task_reschedule_ti_fkey", "task_instance", diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index 72ed418bd0650..089c0da160639 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -60,7 +60,6 @@ func, inspect, or_, - text, tuple_, ) from sqlalchemy.ext.associationproxy import association_proxy @@ -343,7 +342,7 @@ 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, server_default=text("-1")) + map_index = Column(Integer, primary_key=True, nullable=False, default=-1) start_date = Column(UtcDateTime) end_date = Column(UtcDateTime) diff --git a/airflow/models/taskreschedule.py b/airflow/models/taskreschedule.py index 246b0bcd5953c..4fd37a10aad18 100644 --- a/airflow/models/taskreschedule.py +++ b/airflow/models/taskreschedule.py @@ -16,7 +16,7 @@ # specific language governing permissions and limitations # under the License. """TaskReschedule tracks rescheduled task instances.""" -from sqlalchemy import Column, ForeignKeyConstraint, Index, Integer, String, asc, desc, text +from sqlalchemy import Column, ForeignKeyConstraint, Index, Integer, String, asc, desc from sqlalchemy.ext.associationproxy import association_proxy from sqlalchemy.orm import relationship @@ -34,7 +34,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, server_default=text("-1")) + 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) From dcfd72deb85a0fcb0d2e528499b798fc6b90c719 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Sat, 8 Jan 2022 12:21:06 +0800 Subject: [PATCH 07/20] Use StringID in migration --- ...5c0453f75_add_taskmap_and_map_id_on_taskinstance.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) 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 index 119f926d3614a..53d165c81b56f 100644 --- 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 @@ -24,9 +24,9 @@ """ from alembic import op -from sqlalchemy import Column, ForeignKeyConstraint, Integer, String +from sqlalchemy import Column, ForeignKeyConstraint, Integer -from airflow.models.base import COLLATION_ARGS +from airflow.models.base import StringID from airflow.utils.sqlalchemy import ExtendedJSON # Revision identifiers, used by Alembic. @@ -71,9 +71,9 @@ def upgrade(): # Create task_map. op.create_table( "task_map", - Column("dag_id", String(250, **COLLATION_ARGS), primary_key=True), - Column("task_id", String(250, **COLLATION_ARGS), primary_key=True), - Column("run_id", String(250, **COLLATION_ARGS), primary_key=True), + 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), From be4212ad1d36f0e3bf8e432012aed2b97cd2d021 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Sat, 8 Jan 2022 12:27:36 +0800 Subject: [PATCH 08/20] Add map_index as TI and TR init arg --- airflow/models/taskinstance.py | 3 +++ airflow/models/taskreschedule.py | 19 ++++++++++++++++++- 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index 089c0da160639..f10f476111241 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -428,10 +428,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") @@ -764,6 +766,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: diff --git a/airflow/models/taskreschedule.py b/airflow/models/taskreschedule.py index 4fd37a10aad18..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.""" @@ -64,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 From f726cf6b0d43ebb72f1128e8554f2da2b96e31af Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Sat, 8 Jan 2022 16:29:05 +0800 Subject: [PATCH 09/20] Record map_index to refresh_from_db test --- tests/models/test_taskinstance.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py index 4b6a48cd4a21b..12d9776f311a6 100644 --- a/tests/models/test_taskinstance.py +++ b/tests/models/test_taskinstance.py @@ -2060,6 +2060,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 +2086,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." From f1bd8d3e8c2341ab3fd76e2841e432a4f99b0a23 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Sat, 8 Jan 2022 18:57:06 +0800 Subject: [PATCH 10/20] Spellchecker doesn't think "downstreams" is a noun --- airflow/models/baseoperator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/models/baseoperator.py b/airflow/models/baseoperator.py index 7fd5a5136c5fd..da0dbf050c17a 100644 --- a/airflow/models/baseoperator.py +++ b/airflow/models/baseoperator.py @@ -1633,7 +1633,7 @@ def map(self, **kwargs) -> "MappedOperator": return MappedOperator.from_operator(self, kwargs) def has_mapped_dependants(self) -> bool: - """Whether any downstreams depend on this task for mapping.""" + """Whether any downstream dependencies depend on this task for mapping.""" from airflow.utils.task_group import MappedTaskGroup, TaskGroup if not self.has_dag(): From 9e0f71dfad713652cb7055ce94483236e80ad877 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Sat, 8 Jan 2022 19:01:01 +0800 Subject: [PATCH 11/20] Try to make migrations work on MySQL ans MSSQL On MySQL you must not drop an index "needed by a primary key", so we need to drop the primary key first. On both MySQL and MSSQL, a column used by a primary key cannot be nullable (which we already accounted for, but I forgot to add nullable=False in the migration file). --- ...5c0453f75_add_taskmap_and_map_id_on_taskinstance.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) 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 index 53d165c81b56f..91afcb99e6f39 100644 --- 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 @@ -40,16 +40,14 @@ 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_index("idx_task_reschedule_dag_task_run") batch_op.drop_constraint("task_reschedule_ti_fkey", "foreignkey") - if op.get_bind().dialect.name == "mysql": # MySQL also creates an index. - batch_op.drop_index("task_reschedule_ti_fkey") + 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, default=-1)) + 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. @@ -96,10 +94,8 @@ def downgrade(): op.drop_table("task_map") with op.batch_alter_table("task_reschedule") as batch_op: - batch_op.drop_index("idx_task_reschedule_dag_task_run") batch_op.drop_constraint("task_reschedule_ti_fkey", "foreignkey") - if op.get_bind().dialect.name == "mysql": # MySQL also creates an index. - batch_op.drop_index("task_reschedule_ti_fkey") + 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") From 0ab322633b1c708d8b70bfb2c9d252ee50764d2d Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Wed, 12 Jan 2022 16:20:00 +0800 Subject: [PATCH 12/20] Move TaskMap creation logic to TaskMap init --- airflow/models/taskinstance.py | 19 ++----------------- airflow/models/taskmap.py | 34 ++++++++++++++++++++++++++++++++++ 2 files changed, 36 insertions(+), 17 deletions(-) diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index f10f476111241..0ec5d496df859 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -2135,27 +2135,12 @@ def set_duration(self) -> None: self.duration = None self.log.debug("Task Duration set to %s", self.duration) - @provide_session - def _record_task_map_for_downstreams(self, value: Any, *, session: Session = NEW_SESSION) -> None: + 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): return # TODO: Error if the pushed value is not mappable? - session.query(TaskMap).filter_by( - dag_id=self.dag_id, - task_id=self.task_id, - run_id=self.run_id, - map_index=self.map_index, - ).delete() - instance = TaskMap( - dag_id=self.dag_id, - task_id=self.task_id, - run_id=self.run_id, - map_index=self.map_index, - length=len(value), - keys=(list(value) if isinstance(value, collections.abc.Mapping) else None), - ) - session.merge(instance) + session.merge(TaskMap.from_task_instance_xcom(self, value)) @provide_session def xcom_push( diff --git a/airflow/models/taskmap.py b/airflow/models/taskmap.py index 17f7b6f4f589b..20c768b9ac93d 100644 --- a/airflow/models/taskmap.py +++ b/airflow/models/taskmap.py @@ -18,13 +18,18 @@ """Table to store "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. @@ -69,6 +74,35 @@ class TaskMap(Base): ), ) + 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: From d05053dbb72b7ad4771ec58161aec4d6cc1341c5 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Wed, 12 Jan 2022 16:23:47 +0800 Subject: [PATCH 13/20] Fail ti if unmappable value is pushed for mapping --- airflow/models/taskinstance.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index 0ec5d496df859..be1b1358fe9a9 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -109,7 +109,7 @@ from airflow.utils.retries import run_with_db_retries from airflow.utils.session import NEW_SESSION, create_session, provide_session from airflow.utils.sqlalchemy import ExtendedJSON, UtcDateTime -from airflow.utils.state import DagRunState, State +from airflow.utils.state import DagRunState, State, TaskInstanceState from airflow.utils.timeout import timeout try: @@ -2139,7 +2139,10 @@ def _record_task_map_for_downstreams(self, value: Any, *, session: Session) -> N if not self.task.has_mapped_dependants(): return if not isinstance(value, collections.abc.Collection): - return # TODO: Error if the pushed value is not mappable? + self.log.error("Failing %s for unmappable XCom push %r", self.key, value) + self.state = TaskInstanceState.FAILED + session.merge(self) + return session.merge(TaskMap.from_task_instance_xcom(self, value)) @provide_session From 3fe72833823306fe55757e9ff3408474cac22186 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Wed, 12 Jan 2022 21:25:13 +0800 Subject: [PATCH 14/20] Move task map recording call a stack higher This call actually makes more sense in TI._execute_task(), which calls xcom_push() for the return value, since we get to remove a somewhat hackish if check. --- airflow/models/taskinstance.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index be1b1358fe9a9..1d14007fa8604 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -1544,7 +1544,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 @@ -2184,9 +2186,6 @@ def xcom_push( session=session, ) - if key == XCOM_RETURN_KEY: - self._record_task_map_for_downstreams(value, session=session) - @provide_session def xcom_pull( self, From 3c0941f8b3823148b4710e768785962b6db4cca2 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Wed, 12 Jan 2022 21:32:01 +0800 Subject: [PATCH 15/20] Some tests (with a bit mocking) --- airflow/exceptions.py | 11 ++++ airflow/models/taskinstance.py | 11 ++-- tests/conftest.py | 12 +++-- tests/models/test_taskinstance.py | 84 ++++++++++++++++++++++++++++++- 4 files changed, 106 insertions(+), 12 deletions(-) diff --git a/airflow/exceptions.py b/airflow/exceptions.py index d1772153c6ab6..9de5c1f8a725d 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 value {self.value!r}" + + class AirflowDagCycleException(AirflowException): """Raise when there is a cycle in DAG definition.""" diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index 1d14007fa8604..bed607ede8f8f 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -83,6 +83,7 @@ DagRunNotFound, TaskDeferralError, TaskDeferred, + UnmappableXComPushed, ) from airflow.models.base import COLLATION_ARGS, ID_LEN, Base from airflow.models.log import Log @@ -109,7 +110,7 @@ from airflow.utils.retries import run_with_db_retries from airflow.utils.session import NEW_SESSION, create_session, provide_session from airflow.utils.sqlalchemy import ExtendedJSON, UtcDateTime -from airflow.utils.state import DagRunState, State, TaskInstanceState +from airflow.utils.state import DagRunState, State from airflow.utils.timeout import timeout try: @@ -2140,11 +2141,9 @@ def set_duration(self) -> None: 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): - self.log.error("Failing %s for unmappable XCom push %r", self.key, value) - self.state = TaskInstanceState.FAILED - session.merge(self) - 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, value) + raise UnmappableXComPushed(value) session.merge(TaskMap.from_task_instance_xcom(self, value)) @provide_session 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 12d9776f311a6..89eab71494791 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 @@ -2225,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, @@ -2233,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 value 'abc'" + + @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 From 50c2cd76e697955e3791fef3ea2eb32b66627234 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Wed, 12 Jan 2022 21:38:44 +0800 Subject: [PATCH 16/20] Move decorator-mapping logic to MappedOperator This makes it easier to keep the logic in sync with BaseOperator. --- airflow/decorators/base.py | 16 +++++---------- airflow/models/baseoperator.py | 37 ++++++++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 11 deletions(-) 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/models/baseoperator.py b/airflow/models/baseoperator.py index da0dbf050c17a..61b24698d4ac6 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 @@ -1695,6 +1696,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( @@ -1705,9 +1707,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": @@ -1726,8 +1731,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) @@ -1735,6 +1768,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 From 9446b312409538818cbf0cb1bb93d24fa64c5b9c Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Thu, 13 Jan 2022 12:50:50 +0800 Subject: [PATCH 17/20] Do not log the full unmappable XCom value An XCom value is arbitrary and can potentially be a giant blob of data, which would make the log unreadable. Instead, only log the *type* of the XCom value instead; since the offending TaskInstance's key is also logged, the user can always look up the actual value in the XCom storage if needed. The exception still contains a reference to the actual value so the catching frame can access it. Since that value is already in memory anyway, keeping it a bit longer shouldn't be too big a problem. --- airflow/exceptions.py | 2 +- airflow/models/taskinstance.py | 2 +- tests/models/test_taskinstance.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/airflow/exceptions.py b/airflow/exceptions.py index 9de5c1f8a725d..96ad39f4adc55 100644 --- a/airflow/exceptions.py +++ b/airflow/exceptions.py @@ -107,7 +107,7 @@ def __init__(self, value: Any) -> None: self.value = value def __str__(self) -> str: - return f"unmappable return value {self.value!r}" + return f"unmappable return type {type(self.value).__qualname__!r}" class AirflowDagCycleException(AirflowException): diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index bed607ede8f8f..0382ad5ef8792 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -2142,7 +2142,7 @@ def _record_task_map_for_downstreams(self, value: Any, *, session: Session) -> N 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, value) + 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)) diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py index 89eab71494791..e11594f75fcc0 100644 --- a/tests/models/test_taskinstance.py +++ b/tests/models/test_taskinstance.py @@ -2290,7 +2290,7 @@ def push_something(): assert dag_maker.session.query(TaskMap).count() == 0 assert ti.state == TaskInstanceState.FAILED - assert str(ctx.value) == "unmappable return value 'abc'" + assert str(ctx.value) == "unmappable return type 'str'" @pytest.mark.parametrize( "xcom_value, expected_length, expected_keys", From 338a457db18aa1f50cd4803467ca618fd7d9a64c Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Thu, 13 Jan 2022 13:02:53 +0800 Subject: [PATCH 18/20] Note on has_mapped_dependants implementation We need to traverse the entire DAG to find mapped decendants because task groups are not currently recorded in the upstream-downstream lists, only operators. This should change in the future, but until then we need to make do. --- airflow/models/baseoperator.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/airflow/models/baseoperator.py b/airflow/models/baseoperator.py index 61b24698d4ac6..3b9d792223210 100644 --- a/airflow/models/baseoperator.py +++ b/airflow/models/baseoperator.py @@ -1634,7 +1634,13 @@ 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.""" + """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(): From b7038e925e354894ff162d8d28997926f198697a Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Thu, 13 Jan 2022 17:24:07 +0800 Subject: [PATCH 19/20] Add 'unmappable' as a valid word --- docs/spelling_wordlist.txt | 1 + 1 file changed, 1 insertion(+) 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 From 02b6bb44867f06c14d4b1f555c5fde92e75f72d8 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Fri, 14 Jan 2022 13:17:28 +0800 Subject: [PATCH 20/20] Clarify TaskMap costring Co-authored-by: Ash Berlin-Taylor --- airflow/models/taskmap.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/models/taskmap.py b/airflow/models/taskmap.py index 20c768b9ac93d..e91fbd6148c29 100644 --- a/airflow/models/taskmap.py +++ b/airflow/models/taskmap.py @@ -16,7 +16,7 @@ # specific language governing permissions and limitations # under the License. -"""Table to store "mapped" task instances (AIP-42).""" +"""Table to store information about mapped task instances (AIP-42).""" import collections.abc import enum