diff --git a/airflow/api_connexion/openapi/v1.yaml b/airflow/api_connexion/openapi/v1.yaml index 8637612df8194..42c18ab6acd6e 100644 --- a/airflow/api_connexion/openapi/v1.yaml +++ b/airflow/api_connexion/openapi/v1.yaml @@ -3582,6 +3582,13 @@ components: type: string sla_miss: $ref: "#/components/schemas/SLAMiss" + rendered_map_index: + description: | + Rendered name of an expanded task instance, if the task is mapped. + + *New in version 2.9.0* + type: string + nullable: true rendered_fields: description: | JSON object describing rendered fields. diff --git a/airflow/api_connexion/schemas/task_instance_schema.py b/airflow/api_connexion/schemas/task_instance_schema.py index c9084151d88ef..8f0fc0533520d 100644 --- a/airflow/api_connexion/schemas/task_instance_schema.py +++ b/airflow/api_connexion/schemas/task_instance_schema.py @@ -67,6 +67,7 @@ class Meta: executor_config = auto_field() note = auto_field() sla_miss = fields.Nested(SlaMissSchema, dump_default=None) + rendered_map_index = auto_field() rendered_fields = JsonObjectField(dump_default={}) trigger = fields.Nested(TriggerSchema) triggerer_job = fields.Nested(JobSchema) diff --git a/airflow/decorators/base.py b/airflow/decorators/base.py index 51ebbce29c24c..93c403e0bbe1c 100644 --- a/airflow/decorators/base.py +++ b/airflow/decorators/base.py @@ -459,6 +459,7 @@ def _expand(self, expand_input: ExpandInput, *, strict: bool) -> XComArg: expand_input=EXPAND_INPUT_EMPTY, # Don't use this; mapped values go to op_kwargs_expand_input. partial_kwargs=partial_kwargs, task_id=task_id, + map_index_template=partial_kwargs.pop("map_index_template", None), params=partial_params, deps=MappedOperator.deps_for(self.operator_class), operator_extra_links=self.operator_class.operator_extra_links, diff --git a/airflow/models/baseoperator.py b/airflow/models/baseoperator.py index 7e8be8de0332a..18d596bc4ac6b 100644 --- a/airflow/models/baseoperator.py +++ b/airflow/models/baseoperator.py @@ -245,6 +245,7 @@ def partial( priority_weight: int | ArgNotSet = NOTSET, weight_rule: str | ArgNotSet = NOTSET, sla: timedelta | None | ArgNotSet = NOTSET, + map_index_template: str | None | ArgNotSet = NOTSET, max_active_tis_per_dag: int | None | ArgNotSet = NOTSET, max_active_tis_per_dagrun: int | None | ArgNotSet = NOTSET, on_execute_callback: None | TaskStateChangeCallback | list[TaskStateChangeCallback] | ArgNotSet = NOTSET, @@ -289,6 +290,7 @@ def partial( "dag": dag, "task_group": task_group, "task_id": task_id, + "map_index_template": map_index_template, "start_date": start_date, "end_date": end_date, "owner": owner, @@ -781,6 +783,7 @@ def __init__( resources: dict[str, Any] | None = None, run_as_user: str | None = None, task_concurrency: int | None = None, + map_index_template: str | None = None, max_active_tis_per_dag: int | None = None, max_active_tis_per_dagrun: int | None = None, executor_config: dict | None = None, @@ -933,6 +936,7 @@ def __init__( self.max_active_tis_per_dag: int | None = max_active_tis_per_dag self.max_active_tis_per_dagrun: int | None = max_active_tis_per_dagrun self.do_xcom_push: bool = do_xcom_push + self.map_index_template: str | None = map_index_template self.multiple_outputs: bool = multiple_outputs self.doc_md = doc_md @@ -1572,6 +1576,7 @@ def get_serialized_fields(cls): "is_setup", "is_teardown", "on_failure_fail_dagrun", + "map_index_template", } ) DagContext.pop_context_managed_dag() diff --git a/airflow/models/mappedoperator.py b/airflow/models/mappedoperator.py index 3c555e874ff4b..c6f8b5e4569a0 100644 --- a/airflow/models/mappedoperator.py +++ b/airflow/models/mappedoperator.py @@ -211,6 +211,7 @@ def _expand(self, expand_input: ExpandInput, *, strict: bool) -> MappedOperator: expand_input=expand_input, partial_kwargs=partial_kwargs, task_id=task_id, + map_index_template=partial_kwargs.pop("map_index_template", None), params=self.params, deps=MappedOperator.deps_for(self.operator_class), operator_extra_links=self.operator_class.operator_extra_links, @@ -280,6 +281,7 @@ class MappedOperator(AbstractOperator): end_date: pendulum.DateTime | None upstream_task_ids: set[str] = attr.ib(factory=set, init=False) downstream_task_ids: set[str] = attr.ib(factory=set, init=False) + map_index_template: str | None _disallow_kwargs_override: bool """Whether execution fails if ``expand_input`` has duplicates to ``partial_kwargs``. diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index f3f89db86d204..fa0367fb13858 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -750,6 +750,7 @@ def get_triggering_events() -> dict[str, list[DatasetEvent | DatasetEventPydanti "inlets": task.inlets, "logical_date": logical_date, "macros": macros, + "map_index_template": task.map_index_template, "next_ds": get_next_ds(), "next_ds_nodash": get_next_ds_nodash(), "next_execution_date": get_next_execution_date(), @@ -1252,6 +1253,7 @@ class TaskInstance(Base, LoggingMixin): pid = Column(Integer) executor_config = Column(ExecutorConfigType(pickler=dill)) updated_at = Column(UtcDateTime, default=timezone.utcnow, onupdate=timezone.utcnow) + rendered_map_index = Column(String(64)) external_executor_id = Column(StringID()) @@ -2512,7 +2514,12 @@ def signal_handler(signum, frame): self.task.params = context["params"] with set_current_context(context): - task_orig = self.render_templates(context=context) + dag = self.task.get_dag() + if dag is not None: + jinja_env = dag.get_template_env() + else: + jinja_env = None + task_orig = self.render_templates(context=context, jinja_env=jinja_env) if not test_mode: rtif = RenderedTaskInstanceFields(ti=self, render_templates=False) @@ -2547,10 +2554,16 @@ def signal_handler(signum, frame): # Execute the task with set_current_context(context): result = self._execute_task(context, task_orig) + # Run post_execute callback # Is never MappedOperator at this point self.task.post_execute(context=context, result=result) # type: ignore[union-attr] + # DAG authors define map_index_template at the task level + if jinja_env is not None and (template := context.get("map_index_template")) is not None: + rendered_map_index = self.rendered_map_index = jinja_env.from_string(template).render(context) + self.log.info("Map index rendered as %s", rendered_map_index) + Stats.incr(f"operator_successes_{self.task.task_type}", tags=self.stats_tags) # Same metric with tagging Stats.incr("operator_successes", tags={**self.stats_tags, "task_type": self.task.task_type}) @@ -2922,7 +2935,9 @@ def overwrite_params_with_dag_run_conf(self, params: dict, dag_run: DagRun): self.log.debug("Updating task params (%s) with DagRun.conf (%s)", params, dag_run.conf) params.update(dag_run.conf) - def render_templates(self, context: Context | None = None) -> Operator: + def render_templates( + self, context: Context | None = None, jinja_env: jinja2.Environment | None = None + ) -> Operator: """Render templates in the operator fields. If the task was originally mapped, this may replace ``self.task`` with @@ -2937,7 +2952,7 @@ def render_templates(self, context: Context | None = None) -> Operator: # unmapped BaseOperator created by this function! This is because the # MappedOperator is useless for template rendering, and we need to be # able to access the unmapped task instead. - original_task.render_template_fields(context) + original_task.render_template_fields(context, jinja_env) return original_task diff --git a/airflow/operators/python.py b/airflow/operators/python.py index 2fa004f0280a5..da593aea390d6 100644 --- a/airflow/operators/python.py +++ b/airflow/operators/python.py @@ -312,6 +312,7 @@ class _BasePythonVirtualenvOperator(PythonOperator, metaclass=ABCMeta): "ds_nodash", "expanded_ti_count", "inlets", + "map_index_template", "next_ds", "next_ds_nodash", "outlets", diff --git a/airflow/serialization/pydantic/taskinstance.py b/airflow/serialization/pydantic/taskinstance.py index 9e874c569b027..2830618d56d50 100644 --- a/airflow/serialization/pydantic/taskinstance.py +++ b/airflow/serialization/pydantic/taskinstance.py @@ -100,6 +100,7 @@ class TaskInstancePydantic(BaseModelPydantic, LoggingMixin): pid: Optional[int] executor_config: Any updated_at: Optional[datetime] + rendered_map_index: Optional[str] external_executor_id: Optional[str] trigger_id: Optional[int] trigger_timeout: Optional[datetime] diff --git a/airflow/serialization/schema.json b/airflow/serialization/schema.json index 71ee0c80063c0..6a52ac52d971e 100644 --- a/airflow/serialization/schema.json +++ b/airflow/serialization/schema.json @@ -287,7 +287,8 @@ "_log_config_logger_name": { "type": "string" }, "_is_mapped": { "const": true, "$comment": "only present when True" }, "expand_input": { "type": "object" }, - "partial_kwargs": { "type": "object" } + "partial_kwargs": { "type": "object" }, + "map_index_template": { "type": "string" } }, "dependencies": { "expand_input": ["partial_kwargs", "_is_mapped"], diff --git a/airflow/serialization/serialized_objects.py b/airflow/serialization/serialized_objects.py index 552244d73ba7b..0c11fc9301725 100644 --- a/airflow/serialization/serialized_objects.py +++ b/airflow/serialization/serialized_objects.py @@ -1124,6 +1124,7 @@ def deserialize_operator(cls, encoded_op: dict[str, Any]) -> Operator: task_group=None, start_date=None, end_date=None, + map_index_template=None, disallow_kwargs_override=encoded_op["_disallow_kwargs_override"], expand_input_attr=encoded_op["_expand_input_attr"], ) diff --git a/airflow/utils/context.py b/airflow/utils/context.py index e2db6df710fa4..632126062b0b5 100644 --- a/airflow/utils/context.py +++ b/airflow/utils/context.py @@ -59,6 +59,7 @@ "inlets", "logical_date", "macros", + "map_index_template", "next_ds", "next_ds_nodash", "next_execution_date", diff --git a/airflow/utils/context.pyi b/airflow/utils/context.pyi index 9fecccfb1d357..124cd9c8c4ec6 100644 --- a/airflow/utils/context.pyi +++ b/airflow/utils/context.pyi @@ -71,6 +71,7 @@ class Context(TypedDict, total=False): inlets: list logical_date: DateTime macros: Any + map_index_template: str next_ds: str | None next_ds_nodash: str | None next_execution_date: DateTime | None diff --git a/airflow/www/static/js/components/Table/Cells.tsx b/airflow/www/static/js/components/Table/Cells.tsx index 9e86a74935f10..346acea91edd6 100644 --- a/airflow/www/static/js/components/Table/Cells.tsx +++ b/airflow/www/static/js/components/Table/Cells.tsx @@ -39,7 +39,7 @@ import { getMetaValue } from "src/utils"; import { useContainerRef } from "src/context/containerRef"; import { SimpleStatus } from "src/dag/StatusBox"; -interface CellProps { +export interface CellProps { cell: { // eslint-disable-next-line @typescript-eslint/no-explicit-any value: any; diff --git a/airflow/www/static/js/dag/details/Header.tsx b/airflow/www/static/js/dag/details/Header.tsx index a2f5b30eca40c..1903574630862 100644 --- a/airflow/www/static/js/dag/details/Header.tsx +++ b/airflow/www/static/js/dag/details/Header.tsx @@ -28,7 +28,7 @@ import { import { getDagRunLabel, getMetaValue, getTask } from "src/utils"; import useSelection from "src/dag/useSelection"; import Time from "src/components/Time"; -import { useGridData } from "src/api"; +import { useGridData, useTaskInstance } from "src/api"; import RunTypeIcon from "src/components/RunTypeIcon"; import BreadcrumbText from "./BreadcrumbText"; @@ -45,6 +45,15 @@ const Header = () => { onSelect, clearSelection, } = useSelection(); + + const { data: taskInstance } = useTaskInstance({ + dagId, + dagRunId: runId || "", + taskId: taskId || "", + mapIndex, + enabled: mapIndex !== undefined, + }); + const dagRun = dagRuns.find((r) => r.runId === runId); const group = getTask({ taskId, task: groups }); @@ -131,7 +140,10 @@ const Header = () => { - + )} diff --git a/airflow/www/static/js/dag/details/taskInstance/Details.tsx b/airflow/www/static/js/dag/details/taskInstance/Details.tsx index 5e78825e5c289..cc4dea6e9648e 100644 --- a/airflow/www/static/js/dag/details/taskInstance/Details.tsx +++ b/airflow/www/static/js/dag/details/taskInstance/Details.tsx @@ -144,6 +144,12 @@ const Details = ({ gridInstance, taskInstance, group }: Props) => { {taskInstance.mapIndex} )} + {taskInstance?.renderedMapIndex !== undefined && ( + + Rendered Map Index + {taskInstance.renderedMapIndex} + + )} {!!taskInstance?.tryNumber && ( Try Number diff --git a/airflow/www/static/js/dag/details/taskInstance/MappedInstances.tsx b/airflow/www/static/js/dag/details/taskInstance/MappedInstances.tsx index 7b297f788ea91..b236feab68ebf 100644 --- a/airflow/www/static/js/dag/details/taskInstance/MappedInstances.tsx +++ b/airflow/www/static/js/dag/details/taskInstance/MappedInstances.tsx @@ -28,6 +28,7 @@ import { StatusWithNotes } from "src/dag/StatusBox"; import { Table } from "src/components/Table"; import Time from "src/components/Time"; import { useOffsetTop } from "src/utils"; +import type { CellProps } from "src/components/Table"; interface Props { dagId: string; @@ -69,6 +70,7 @@ const MappedInstances = ({ dagId, runId, taskId, onRowClicked }: Props) => { () => taskInstances.map((mi) => ({ ...mi, + renderedMapIndex: mi.renderedMapIndex, state: ( { { Header: "Map Index", accessor: "mapIndex", + Cell: ({ cell: { row } }: CellProps) => + row.original.renderedMapIndex || row.original.mapIndex, }, { Header: "State", diff --git a/airflow/www/static/js/types/api-generated.ts b/airflow/www/static/js/types/api-generated.ts index 1184dd5323732..7430dba2724b4 100644 --- a/airflow/www/static/js/types/api-generated.ts +++ b/airflow/www/static/js/types/api-generated.ts @@ -1443,6 +1443,12 @@ export interface components { pid?: number | null; executor_config?: string; sla_miss?: components["schemas"]["SLAMiss"]; + /** + * @description Rendered name of an expanded task instance, if the task is mapped. + * + * *New in version 2.9.0* + */ + rendered_map_index?: string | null; /** * @description JSON object describing rendered fields. * diff --git a/docs/apache-airflow/authoring-and-scheduling/dynamic-task-mapping.rst b/docs/apache-airflow/authoring-and-scheduling/dynamic-task-mapping.rst index 81102dd54e6e2..acf9cfba3dcbb 100644 --- a/docs/apache-airflow/authoring-and-scheduling/dynamic-task-mapping.rst +++ b/docs/apache-airflow/authoring-and-scheduling/dynamic-task-mapping.rst @@ -171,6 +171,44 @@ As well as a single parameter it is possible to pass multiple parameters to expa This would result in the add task being called 6 times. Please note, however, that the order of expansion is not guaranteed. +Named mapping +------------- + +By default, mapped tasks are assigned an integer index. It is possible to override the integer index for each mapped task in the Airflow UI with a name based on the task's input. This is done by providing a Jinja template for the task with ``map_index_template``. This template is rendered after each expanded task is executed using the task context. This means you can reference attributes on the task like this: + +.. code-block:: python + + from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator + + + # The two expanded task instances will be named "2024-01-01" and "2024-01-02". + SQLExecuteQueryOperator.partial( + ..., + sql="SELECT * FROM data WHERE date = %(date)s", + map_index_template="""{{ task.parameters['date'] }}""", + ).expand( + parameters=[{"date": "2024-01-01"}, {"date": "2024-01-02"}], + ) + +In the above example, the expanded task instances will be named "2024-01-01" and "2024-01-02". The names show up in the Airflow UI instead of "0" and "1", respectively. + +Since the template is rendered after the main execution block, it is possible to also dynamically inject into the rendering context. This is useful when the logic to render a desirable name is difficult to express in the Jinja template syntax, particularly in a taskflow function. For example: + +.. code-block:: python + + from airflow.operators.python import get_current_context + + + @task(map_index_template="{{ my_variable }}") + def my_task(my_value: str): + context = get_current_context() + context["my_variable"] = my_value * 3 + ... # Normal execution... + + + # The task instances will be named "aaa" and "bbb". + my_task.expand(my_value=["a", "b"]) + Mapping with non-TaskFlow operators =================================== diff --git a/docs/apache-airflow/templates-ref.rst b/docs/apache-airflow/templates-ref.rst index 9184bdc770c67..0929c8b0c08fa 100644 --- a/docs/apache-airflow/templates-ref.rst +++ b/docs/apache-airflow/templates-ref.rst @@ -74,6 +74,7 @@ Variable Type Description ``{{ run_id }}`` str The currently running :class:`~airflow.models.dagrun.DagRun` run ID. ``{{ dag_run }}`` DagRun The currently running :class:`~airflow.models.dagrun.DagRun`. ``{{ test_mode }}`` bool Whether the task instance was run by the ``airflow test`` CLI. +``{{ map_index_template }}`` None | str Template used to render the expanded task instance of a mapped task. Setting this value will be reflected in the rendered result. ``{{ expanded_ti_count }}`` int | ``None`` | Number of task instances that a mapped task was expanded into. If | the current task is not mapped, this should be ``None``. | Added in version 2.5. diff --git a/tests/api_connexion/endpoints/test_mapped_task_instance_endpoint.py b/tests/api_connexion/endpoints/test_mapped_task_instance_endpoint.py index b7d88e347fc72..30c227ea9a15d 100644 --- a/tests/api_connexion/endpoints/test_mapped_task_instance_endpoint.py +++ b/tests/api_connexion/endpoints/test_mapped_task_instance_endpoint.py @@ -234,6 +234,7 @@ def test_mapped_task_instances(self, one_task_with_mapped_tis, session): "queue": "default", "queued_when": None, "rendered_fields": {}, + "rendered_map_index": None, "sla_miss": None, "start_date": "2020-01-01T00:00:00+00:00", "state": "success", diff --git a/tests/api_connexion/endpoints/test_task_instance_endpoint.py b/tests/api_connexion/endpoints/test_task_instance_endpoint.py index 16849d6cfa6d3..d85751a8c97c5 100644 --- a/tests/api_connexion/endpoints/test_task_instance_endpoint.py +++ b/tests/api_connexion/endpoints/test_task_instance_endpoint.py @@ -245,6 +245,7 @@ def test_should_respond_200(self, username, session): "unixname": getuser(), "dag_run_id": "TEST_DAG_RUN_ID", "rendered_fields": {}, + "rendered_map_index": None, "trigger": None, "triggerer_job": None, } @@ -301,6 +302,7 @@ def test_should_respond_200_with_task_state_in_deferred(self, session): "unixname": getuser(), "dag_run_id": "TEST_DAG_RUN_ID", "rendered_fields": {}, + "rendered_map_index": None, "trigger": { "classpath": "none", "kwargs": "{}", @@ -346,6 +348,7 @@ def test_should_respond_200_with_task_state_in_removed(self, session): "unixname": getuser(), "dag_run_id": "TEST_DAG_RUN_ID", "rendered_fields": {}, + "rendered_map_index": None, "trigger": None, "triggerer_job": None, } @@ -402,6 +405,7 @@ def test_should_respond_200_task_instance_with_sla_and_rendered(self, session): "unixname": getuser(), "dag_run_id": "TEST_DAG_RUN_ID", "rendered_fields": {"op_args": [], "op_kwargs": {}, "templates_dict": None}, + "rendered_map_index": None, "trigger": None, "triggerer_job": None, } @@ -452,6 +456,7 @@ def test_should_respond_200_mapped_task_instance_with_rtif(self, session): "unixname": getuser(), "dag_run_id": "TEST_DAG_RUN_ID", "rendered_fields": {"op_args": [], "op_kwargs": {}, "templates_dict": None}, + "rendered_map_index": None, "trigger": None, "triggerer_job": None, } @@ -2355,6 +2360,7 @@ def test_should_respond_200(self, session): "unixname": getuser(), "dag_run_id": "TEST_DAG_RUN_ID", "rendered_fields": {}, + "rendered_map_index": None, "trigger": None, "triggerer_job": None, } @@ -2409,6 +2415,7 @@ def test_should_respond_200_mapped_task_instance_with_rtif(self, session): "unixname": getuser(), "dag_run_id": "TEST_DAG_RUN_ID", "rendered_fields": {"op_args": [], "op_kwargs": {}, "templates_dict": None}, + "rendered_map_index": None, "trigger": None, "triggerer_job": None, } diff --git a/tests/api_connexion/schemas/test_task_instance_schema.py b/tests/api_connexion/schemas/test_task_instance_schema.py index 55acbb8b958a0..ad2b49a8606ac 100644 --- a/tests/api_connexion/schemas/test_task_instance_schema.py +++ b/tests/api_connexion/schemas/test_task_instance_schema.py @@ -93,6 +93,7 @@ def test_task_instance_schema_without_sla_and_rendered(self, session): "unixname": getuser(), "dag_run_id": None, "rendered_fields": {}, + "rendered_map_index": None, "trigger": None, "triggerer_job": None, } @@ -146,6 +147,7 @@ def test_task_instance_schema_with_sla_and_rendered(self, session): "unixname": getuser(), "dag_run_id": None, "rendered_fields": {"partitions": "data/ds=2022-02-17"}, + "rendered_map_index": None, "trigger": None, "triggerer_job": None, } diff --git a/tests/models/test_mappedoperator.py b/tests/models/test_mappedoperator.py index 253aab8eff70c..5bd7d9a0c1062 100644 --- a/tests/models/test_mappedoperator.py +++ b/tests/models/test_mappedoperator.py @@ -25,6 +25,7 @@ import pendulum import pytest +from sqlalchemy import select from airflow.decorators import setup, task, task_group, teardown from airflow.exceptions import AirflowSkipException @@ -609,6 +610,77 @@ def test_expand_kwargs_mapped_task_instance(dag_maker, session, num_existing_tis assert indices == expected +def _create_mapped_with_name_template_classic(*, task_id, map_names, template): + class HasMapName(BaseOperator): + def __init__(self, *, map_name: str, **kwargs): + super().__init__(**kwargs) + self.map_name = map_name + + def execute(self, context): + context["map_name"] = self.map_name + + return HasMapName.partial(task_id=task_id, map_index_template=template).expand( + map_name=map_names, + ) + + +def _create_mapped_with_name_template_taskflow(*, task_id, map_names, template): + from airflow.operators.python import get_current_context + + @task(task_id=task_id, map_index_template=template) + def task1(map_name): + context = get_current_context() + context["map_name"] = map_name + + return task1.expand(map_name=map_names) + + +@pytest.mark.parametrize( + "template, expected_rendered_names", + [ + pytest.param(None, [None, None], id="unset"), + pytest.param("", ["", ""], id="constant"), + pytest.param("{{ ti.task_id }}-{{ ti.map_index }}", ["task1-0", "task1-1"], id="builtin"), + pytest.param("{{ ti.task_id }}-{{ map_name }}", ["task1-a", "task1-b"], id="custom"), + ], +) +@pytest.mark.parametrize( + "create_mapped_task", + [ + pytest.param(_create_mapped_with_name_template_classic, id="classic"), + pytest.param(_create_mapped_with_name_template_taskflow, id="taskflow"), + ], +) +def test_expand_mapped_task_instance_with_named_index( + dag_maker, + session, + create_mapped_task, + template, + expected_rendered_names, +) -> None: + """Test that the correct number of downstream tasks are generated when mapping with an XComArg""" + with dag_maker("test-dag", session=session, start_date=DEFAULT_DATE): + create_mapped_task(task_id="task1", map_names=["a", "b"], template=template) + + dr = dag_maker.create_dagrun() + tis = dr.get_task_instances() + for ti in tis: + ti.run() + session.flush() + + indices = session.scalars( + select(TaskInstance.rendered_map_index) + .where( + TaskInstance.dag_id == "test-dag", + TaskInstance.task_id == "task1", + TaskInstance.run_id == dr.run_id, + ) + .order_by(TaskInstance.map_index) + ).all() + + assert indices == expected_rendered_names + + @pytest.mark.parametrize( "map_index, expected", [ diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py index 8ed5b7ef99aea..d820c7d3911e2 100644 --- a/tests/models/test_taskinstance.py +++ b/tests/models/test_taskinstance.py @@ -3190,6 +3190,7 @@ def test_refresh_from_db(self, create_task_instance): "operator": "some_custom_operator", "custom_operator_name": "some_custom_operator", "queued_dttm": run_date + datetime.timedelta(hours=1), + "rendered_map_index": None, "queued_by_job_id": 321, "pid": 123, "executor_config": {"Some": {"extra": "information"}}, diff --git a/tests/serialization/test_dag_serialization.py b/tests/serialization/test_dag_serialization.py index 2adc956b6f4d0..63abaa15bc076 100644 --- a/tests/serialization/test_dag_serialization.py +++ b/tests/serialization/test_dag_serialization.py @@ -619,6 +619,7 @@ def validate_deserialized_task( assert isinstance(serialized_task, MappedOperator) fields_to_check = {f.name for f in attr.fields(MappedOperator)} fields_to_check -= { + "map_index_template", # Matching logic in BaseOperator.get_serialized_fields(). "dag", "task_group", @@ -1236,6 +1237,7 @@ def test_no_new_fields_added_to_base_operator(self): "executor_config": {}, "ignore_first_depends_on_past": True, "inlets": [], + "map_index_template": None, "max_active_tis_per_dag": None, "max_active_tis_per_dagrun": None, "max_retry_delay": None, diff --git a/tests/www/views/test_views_tasks.py b/tests/www/views/test_views_tasks.py index 00803a3575e35..da455e4e6015f 100644 --- a/tests/www/views/test_views_tasks.py +++ b/tests/www/views/test_views_tasks.py @@ -1107,6 +1107,7 @@ def test_task_instances(admin_client): "queue": "default", "queued_by_job_id": None, "queued_dttm": None, + "rendered_map_index": None, "run_id": "TEST_DAGRUN", "start_date": None, "state": None, @@ -1139,6 +1140,7 @@ def test_task_instances(admin_client): "queue": "default", "queued_by_job_id": None, "queued_dttm": None, + "rendered_map_index": None, "run_id": "TEST_DAGRUN", "start_date": None, "state": None, @@ -1171,6 +1173,7 @@ def test_task_instances(admin_client): "queue": "default", "queued_by_job_id": None, "queued_dttm": None, + "rendered_map_index": None, "run_id": "TEST_DAGRUN", "start_date": None, "state": None, @@ -1203,6 +1206,7 @@ def test_task_instances(admin_client): "queue": "default", "queued_by_job_id": None, "queued_dttm": None, + "rendered_map_index": None, "run_id": "TEST_DAGRUN", "start_date": None, "state": None, @@ -1235,6 +1239,7 @@ def test_task_instances(admin_client): "queue": "default", "queued_by_job_id": None, "queued_dttm": None, + "rendered_map_index": None, "run_id": "TEST_DAGRUN", "start_date": None, "state": None, @@ -1267,6 +1272,7 @@ def test_task_instances(admin_client): "queue": "default", "queued_by_job_id": None, "queued_dttm": None, + "rendered_map_index": None, "run_id": "TEST_DAGRUN", "start_date": None, "state": None, @@ -1299,6 +1305,7 @@ def test_task_instances(admin_client): "queue": "default", "queued_by_job_id": None, "queued_dttm": None, + "rendered_map_index": None, "run_id": "TEST_DAGRUN", "start_date": None, "state": None,