Skip to content

Commit d08284e

Browse files
uranusjruranusjr
authored
Add map_index to XCom model and interface (#22112)
* Add map_index to XCom primary key This is not actually stored correctly yet. We still need to fix the XCom interface. * Add map_index to XCom interface This adds an additional (optional) map_index argument to XCom's get/set/clear interface so mapped task instances can push to the correct entries, and have them pulled correctly by a downstream. To make the XCom interface easier to use for common scenarios, a convenience method get_value is added to take a TaskInstanceKey that automatically performs argument unpacking and call get_one underneath. This is not done as a get_one overload to simplify the implementation and typing.
1 parent 9eb1a1c commit d08284e

16 files changed

Lines changed: 167 additions & 77 deletions

File tree

‎airflow/migrations/versions/0102_c306b5b5ae4a_switch_xcom_table_to_use_run_id.py‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
# specific language governing permissions and limitations
1717
# under the License.
1818

19-
"""Switch XCom table to use ``run_id``.
19+
"""Switch XCom table to use ``run_id`` and add ``map_index``.
2020
2121
Revision ID: c306b5b5ae4a
2222
Revises: a3bcd0914482
@@ -25,7 +25,7 @@
2525
from typing import Sequence
2626

2727
from alembic import op
28-
from sqlalchemy import Column, Integer, LargeBinary, MetaData, Table, and_, select
28+
from sqlalchemy import Column, Integer, LargeBinary, MetaData, Table, and_, literal_column, select
2929

3030
from airflow.migrations.db_types import TIMESTAMP, StringID
3131
from airflow.migrations.utils import get_mssql_table_constraints
@@ -50,6 +50,7 @@ def _get_new_xcom_columns() -> Sequence[Column]:
5050
Column("timestamp", TIMESTAMP, nullable=False),
5151
Column("dag_id", StringID(), nullable=False),
5252
Column("run_id", StringID(), nullable=False),
53+
Column("map_index", Integer, nullable=False, server_default="-1"),
5354
]
5455

5556

@@ -98,6 +99,7 @@ def upgrade():
9899
xcom.c.timestamp,
99100
xcom.c.dag_id,
100101
dagrun.c.run_id,
102+
literal_column("-1"),
101103
],
102104
).select_from(
103105
xcom.join(
@@ -118,9 +120,9 @@ def upgrade():
118120
op.rename_table("__airflow_tmp_xcom", "xcom")
119121

120122
with op.batch_alter_table("xcom") as batch_op:
121-
batch_op.create_primary_key("xcom_pkey", ["dag_run_id", "task_id", "key"])
123+
batch_op.create_primary_key("xcom_pkey", ["dag_run_id", "task_id", "map_index", "key"])
122124
batch_op.create_index("idx_xcom_key", ["key"])
123-
batch_op.create_index("idx_xcom_ti_id", ["dag_id", "task_id", "run_id"])
125+
batch_op.create_index("idx_xcom_ti_id", ["dag_id", "run_id", "task_id", "map_index"])
124126

125127

126128
def downgrade():
@@ -132,6 +134,10 @@ def downgrade():
132134
op.create_table("__airflow_tmp_xcom", *_get_old_xcom_columns())
133135

134136
xcom = Table("xcom", metadata, *_get_new_xcom_columns())
137+
138+
# Remoe XCom entries from mapped tis.
139+
op.execute(xcom.delete().where(xcom.c.map_index != -1))
140+
135141
dagrun = _get_dagrun_table()
136142
query = select(
137143
[

‎airflow/models/xcom.py‎

Lines changed: 101 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@ class BaseXCom(Base, LoggingMixin):
6161

6262
dag_run_id = Column(Integer(), nullable=False, primary_key=True)
6363
task_id = Column(String(ID_LEN, **COLLATION_ARGS), nullable=False, primary_key=True)
64+
map_index = Column(Integer, primary_key=True, nullable=False, server_default="-1")
6465
key = Column(String(512, **COLLATION_ARGS), nullable=False, primary_key=True)
6566

6667
# Denormalized for easier lookup.
@@ -87,7 +88,7 @@ class BaseXCom(Base, LoggingMixin):
8788
# but it goes over MySQL's index length limit. So we instead create indexes
8889
# separately, and enforce uniqueness with DagRun.id instead.
8990
Index("idx_xcom_key", key),
90-
Index("idx_xcom_ti_id", dag_id, task_id, run_id),
91+
Index("idx_xcom_ti_id", dag_id, task_id, run_id, map_index),
9192
)
9293

9394
@reconstructor
@@ -111,6 +112,7 @@ def set(
111112
dag_id: str,
112113
task_id: str,
113114
run_id: str,
115+
map_index: int = -1,
114116
session: Session = NEW_SESSION,
115117
) -> None:
116118
"""Store an XCom value.
@@ -123,6 +125,8 @@ def set(
123125
:param dag_id: DAG ID.
124126
:param task_id: Task ID.
125127
:param run_id: DAG run ID for the task.
128+
:param map_index: Optional map index to assign XCom for a mapped task.
129+
The default is ``-1`` (set for a non-mapped task).
126130
:param session: Database session. If not given, a new session will be
127131
created for this function.
128132
"""
@@ -152,6 +156,7 @@ def set(
152156
session: Session = NEW_SESSION,
153157
*,
154158
run_id: Optional[str] = None,
159+
map_index: int = -1,
155160
) -> None:
156161
""":sphinx-autoapi-skip:"""
157162
from airflow.models.dagrun import DagRun
@@ -183,6 +188,7 @@ def set(
183188
task_id=task_id,
184189
dag_id=dag_id,
185190
run_id=run_id,
191+
map_index=map_index,
186192
)
187193

188194
# Remove duplicate XComs and insert a new one.
@@ -203,13 +209,49 @@ def set(
203209
session.add(new)
204210
session.flush()
205211

212+
@classmethod
213+
@provide_session
214+
def get_value(
215+
cls,
216+
*,
217+
ti_key: "TaskInstanceKey",
218+
key: Optional[str] = None,
219+
session: Session = NEW_SESSION,
220+
) -> Any:
221+
"""Retrieve an XCom value for a task instance.
222+
223+
This method returns "full" XCom values (i.e. uses ``deserialize_value``
224+
from the XCom backend). Use :meth:`get_many` if you want the "shortened"
225+
value via ``orm_deserialize_value``.
226+
227+
If there are no results, *None* is returned. If multiple XCom entries
228+
match the criteria, an arbitrary one is returned.
229+
230+
:param ti_key: The TaskInstanceKey to look up the XCom for.
231+
:param key: A key for the XCom. If provided, only XCom with matching
232+
keys will be returned. Pass *None* (default) to remove the filter.
233+
:param session: Database session. If not given, a new session will be
234+
created for this function.
235+
"""
236+
return cls.get_one(
237+
key=key,
238+
task_id=ti_key.task_id,
239+
dag_id=ti_key.dag_id,
240+
run_id=ti_key.run_id,
241+
map_index=ti_key.map_index,
242+
session=session,
243+
)
244+
206245
@overload
207246
@classmethod
208247
def get_one(
209248
cls,
210249
*,
211250
key: Optional[str] = None,
212-
ti_key: "TaskInstanceKey",
251+
dag_id: Optional[str] = None,
252+
task_id: Optional[str] = None,
253+
run_id: Optional[str] = None,
254+
map_index: Optional[int] = None,
213255
session: Session = NEW_SESSION,
214256
) -> Optional[Any]:
215257
"""Retrieve an XCom value, optionally meeting certain criteria.
@@ -218,12 +260,22 @@ def get_one(
218260
from the XCom backend). Use :meth:`get_many` if you want the "shortened"
219261
value via ``orm_deserialize_value``.
220262
221-
If there are no results, *None* is returned.
263+
If there are no results, *None* is returned. If multiple XCom entries
264+
match the criteria, an arbitrary one is returned.
222265
223266
A deprecated form of this function accepts ``execution_date`` instead of
224267
``run_id``. The two arguments are mutually exclusive.
225268
226-
:param ti_key: The TaskInstanceKey to look up the XCom for
269+
.. seealso:: ``get_value()`` is a convenience function if you already
270+
have a structured TaskInstance or TaskInstanceKey object available.
271+
272+
:param run_id: DAG run ID for the task.
273+
:param dag_id: Only pull XCom from this DAG. Pass *None* (default) to
274+
remove the filter.
275+
:param task_id: Only XCom from task with matching ID will be pulled.
276+
Pass *None* (default) to remove the filter.
277+
:param map_index: Only XCom from task with matching ID will be pulled.
278+
Pass *None* (default) to remove the filter.
227279
:param key: A key for the XCom. If provided, only XCom with matching
228280
keys will be returned. Pass *None* (default) to remove the filter.
229281
:param include_prior_dates: If *False* (default), only XCom from the
@@ -233,19 +285,6 @@ def get_one(
233285
created for this function.
234286
"""
235287

236-
@overload
237-
@classmethod
238-
def get_one(
239-
cls,
240-
*,
241-
key: Optional[str] = None,
242-
task_id: str,
243-
dag_id: str,
244-
run_id: str,
245-
session: Session = NEW_SESSION,
246-
) -> Optional[Any]:
247-
...
248-
249288
@overload
250289
@classmethod
251290
def get_one(
@@ -271,27 +310,19 @@ def get_one(
271310
session: Session = NEW_SESSION,
272311
*,
273312
run_id: Optional[str] = None,
274-
ti_key: Optional["TaskInstanceKey"] = None,
313+
map_index: Optional[int] = None,
275314
) -> Optional[Any]:
276315
""":sphinx-autoapi-skip:"""
277-
if not exactly_one(execution_date is not None, ti_key is not None, run_id is not None):
316+
if not exactly_one(execution_date is not None, run_id is not None):
278317
raise ValueError("Exactly one of ti_key, run_id, or execution_date must be passed")
279318

280-
if ti_key is not None:
281-
query = session.query(cls).filter_by(
282-
dag_id=ti_key.dag_id,
283-
run_id=ti_key.run_id,
284-
task_id=ti_key.task_id,
285-
)
286-
if key:
287-
query = query.filter_by(key=key)
288-
query = query.limit(1)
289-
elif run_id:
319+
if run_id:
290320
query = cls.get_many(
291321
run_id=run_id,
292322
key=key,
293323
task_ids=task_id,
294324
dag_ids=dag_id,
325+
map_indexes=map_index,
295326
include_prior_dates=include_prior_dates,
296327
limit=1,
297328
session=session,
@@ -307,6 +338,7 @@ def get_one(
307338
key=key,
308339
task_ids=task_id,
309340
dag_ids=dag_id,
341+
map_indexes=map_index,
310342
include_prior_dates=include_prior_dates,
311343
limit=1,
312344
session=session,
@@ -328,6 +360,7 @@ def get_many(
328360
key: Optional[str] = None,
329361
task_ids: Union[str, Iterable[str], None] = None,
330362
dag_ids: Union[str, Iterable[str], None] = None,
363+
map_indexes: Union[int, Iterable[int], None] = None,
331364
include_prior_dates: bool = False,
332365
limit: Optional[int] = None,
333366
session: Session = NEW_SESSION,
@@ -347,6 +380,8 @@ def get_many(
347380
Pass *None* (default) to remove the filter.
348381
:param dag_id: Only pulls XComs from this DAG. If *None* (default), the
349382
DAG of the calling task is used.
383+
:param map_index: Only XComs from matching map indexes will be pulled.
384+
Pass *None* (default) to remove the filter.
350385
:param include_prior_dates: If *False* (default), only XComs from the
351386
specified DAG run are returned. If *True*, all matching XComs are
352387
returned regardless of the run it belongs to.
@@ -362,6 +397,7 @@ def get_many(
362397
key: Optional[str] = None,
363398
task_ids: Union[str, Iterable[str], None] = None,
364399
dag_ids: Union[str, Iterable[str], None] = None,
400+
map_indexes: Union[int, Iterable[int], None] = None,
365401
include_prior_dates: bool = False,
366402
limit: Optional[int] = None,
367403
session: Session = NEW_SESSION,
@@ -376,6 +412,7 @@ def get_many(
376412
key: Optional[str] = None,
377413
task_ids: Optional[Union[str, Iterable[str]]] = None,
378414
dag_ids: Optional[Union[str, Iterable[str]]] = None,
415+
map_indexes: Union[int, Iterable[int], None] = None,
379416
include_prior_dates: bool = False,
380417
limit: Optional[int] = None,
381418
session: Session = NEW_SESSION,
@@ -406,6 +443,11 @@ def get_many(
406443
elif dag_ids is not None:
407444
query = query.filter(cls.dag_id == dag_ids)
408445

446+
if is_container(map_indexes):
447+
query = query.filter(cls.map_index.in_(map_indexes))
448+
elif map_indexes is not None:
449+
query = query.filter(cls.map_index == map_indexes)
450+
409451
if include_prior_dates:
410452
if execution_date is not None:
411453
query = query.filter(DagRun.execution_date <= execution_date)
@@ -438,7 +480,15 @@ def delete(cls, xcoms: Union["XCom", Iterable["XCom"]], session: Session) -> Non
438480

439481
@overload
440482
@classmethod
441-
def clear(cls, *, dag_id: str, task_id: str, run_id: str, session: Optional[Session] = None) -> None:
483+
def clear(
484+
cls,
485+
*,
486+
dag_id: str,
487+
task_id: str,
488+
run_id: str,
489+
map_index: Optional[int] = None,
490+
session: Session = NEW_SESSION,
491+
) -> None:
442492
"""Clear all XCom data from the database for the given task instance.
443493
444494
A deprecated form of this function accepts ``execution_date`` instead of
@@ -447,6 +497,8 @@ def clear(cls, *, dag_id: str, task_id: str, run_id: str, session: Optional[Sess
447497
:param dag_id: ID of DAG to clear the XCom for.
448498
:param task_id: ID of task to clear the XCom for.
449499
:param run_id: ID of DAG run to clear the XCom for.
500+
:param map_index: If given, only clear XCom from this particular mapped
501+
task. The default ``None`` clears *all* XComs from the task.
450502
:param session: Database session. If not given, a new session will be
451503
created for this function.
452504
"""
@@ -472,6 +524,7 @@ def clear(
472524
session: Session = NEW_SESSION,
473525
*,
474526
run_id: Optional[str] = None,
527+
map_index: Optional[int] = None,
475528
) -> None:
476529
""":sphinx-autoapi-skip:"""
477530
from airflow.models import DagRun
@@ -495,17 +548,20 @@ def clear(
495548
.scalar()
496549
)
497550

498-
return session.query(cls).filter_by(dag_id=dag_id, task_id=task_id, run_id=run_id).delete()
551+
query = session.query(cls).filter_by(dag_id=dag_id, task_id=task_id, run_id=run_id)
552+
if map_index is not None:
553+
query = query.filter_by(map_index=map_index)
554+
query.delete()
499555

500556
@staticmethod
501557
def serialize_value(
502558
value: Any,
503559
*,
504-
key=None,
505-
task_id=None,
506-
dag_id=None,
507-
run_id=None,
508-
mapping_index: int = -1,
560+
key: Optional[str] = None,
561+
task_id: Optional[str] = None,
562+
dag_id: Optional[str] = None,
563+
run_id: Optional[str] = None,
564+
map_index: Optional[int] = None,
509565
):
510566
"""Serialize XCom value to str or pickled object"""
511567
if conf.getboolean('core', 'enable_xcom_pickling'):
@@ -549,13 +605,14 @@ def orm_deserialize_value(self) -> Any:
549605
return BaseXCom.deserialize_value(self)
550606

551607

552-
def _patch_outdated_serializer(clazz, params):
553-
"""
554-
Previously XCom.serialize_value only accepted one argument ``value``. In order to give
555-
custom XCom backends more flexibility with how they store values we now forward to
556-
``XCom.serialize_value`` all params passed to ``XCom.set``. In order to maintain
557-
compatibility with XCom backends written with the old signature we check the signature
558-
and if necessary we patch with a method that ignores kwargs the backend does not accept.
608+
def _patch_outdated_serializer(clazz: Type[BaseXCom], params: Iterable[str]) -> None:
609+
"""Patch a custom ``serialize_value`` to accept the modern signature.
610+
611+
To give custom XCom backends more flexibility with how they store values, we
612+
now forward all params passed to ``XCom.set`` to ``XCom.serialize_value``.
613+
In order to maintain compatibility with custom XCom backends written with
614+
the old signature, we check the signature and, if necessary, patch with a
615+
method that ignores kwargs the backend does not accept.
559616
"""
560617
old_serializer = clazz.serialize_value
561618

@@ -570,7 +627,7 @@ def _shim(**kwargs):
570627
)
571628
return old_serializer(**kwargs)
572629

573-
clazz.serialize_value = _shim
630+
clazz.serialize_value = _shim # type: ignore[assignment]
574631

575632

576633
def _get_function_params(function) -> List[str]:

‎airflow/operators/trigger_dagrun.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ def get_link(
5656
) -> str:
5757
# Fetch the correct execution date for the triggerED dag which is
5858
# stored in xcom during execution of the triggerING task.
59-
when = XCom.get_one(ti_key=ti_key, key=XCOM_EXECUTION_DATE_ISO)
59+
when = XCom.get_value(ti_key=ti_key, key=XCOM_EXECUTION_DATE_ISO)
6060
query = {"dag_id": cast(TriggerDagRunOperator, operator).trigger_dag_id, "base_date": when}
6161
return build_airflow_url_with_query(query)
6262

‎airflow/providers/amazon/aws/operators/emr.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -245,7 +245,7 @@ def get_link(
245245
:return: url link
246246
"""
247247
if ti_key:
248-
flow_id = XCom.get_one(key="return_value", ti_key=ti_key)
248+
flow_id = XCom.get_value(key="return_value", ti_key=ti_key)
249249
else:
250250
assert dttm
251251
flow_id = XCom.get_one(

0 commit comments

Comments
 (0)