This is an automated email from the ASF dual-hosted git repository.
ashb pushed a commit to branch task-loops-stack-4
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/task-loops-stack-4 by this
push:
new 8b8d40924b1 fixup! Pin loop definitions and carry region coordinates
through task execution
8b8d40924b1 is described below
commit 8b8d40924b125059ef0fda024fd1e607eed86b7b
Author: Ash Berlin-Taylor <[email protected]>
AuthorDate: Wed Oct 7 23:58:21 2026 +0100
fixup! Pin loop definitions and carry region coordinates through task
execution
Co-Authored-By: Claude <[email protected]>
---
.../execution_api/routes/task_instances.py | 223 ++++++++++--------
.../api_fastapi/execution_api/routes/xcoms.py | 67 +-----
.../api_fastapi/execution_api/versions/__init__.py | 4 -
.../versions/v2026_10_30_xcom_params.py | 37 ---
.../src/airflow/executors/workloads/task.py | 3 +-
.../src/airflow/models/task_coordinates.py | 93 +++++++-
.../src/airflow/utils/log/task_log_address.py | 25 ++-
.../versions/head/test_task_instances.py | 174 ++++++++++----
.../execution_api/versions/head/test_xcoms.py | 34 +--
.../tests/unit/executors/test_workloads.py | 2 +
.../tests/unit/models/test_task_coordinates.py | 28 ++-
.../tests/unit/utils/log/test_task_log_address.py | 35 +++
go-sdk/airflow/spec_fields.gen.go | 1 -
go-sdk/pkg/execution/genmodels/models.gen.go | 106 ++-------
go-sdk/schema/supervisor-schema.json | 250 ---------------------
.../cncf/kubernetes/template_rendering.py | 2 +-
.../cncf/kubernetes/test_template_rendering.py | 4 +-
.../providers/edge3/worker_api/routes/logs.py | 18 ++
.../edge3/worker_api/v2-edge-generated.yaml | 11 +
.../unit/edge3/worker_api/routes/test_logs.py | 43 ++++
.../providers/elasticsearch/log/es_task_handler.py | 3 +-
.../unit/elasticsearch/log/test_es_task_handler.py | 6 +
.../providers/opensearch/log/os_task_handler.py | 3 +-
.../unit/opensearch/log/test_os_task_handler.py | 6 +
task-sdk/src/airflow/sdk/api/client.py | 60 +----
task-sdk/src/airflow/sdk/execution_time/comms.py | 20 --
.../airflow/sdk/execution_time/request_handlers.py | 20 +-
.../airflow/sdk/execution_time/schema/schema.json | 250 ---------------------
.../sdk/execution_time/schema/versions/__init__.py | 2 -
.../execution_time/schema/versions/v2026_10_30.py | 38 ----
.../src/airflow/sdk/execution_time/supervisor.py | 2 -
.../execution_time/schema/test_migrator.py | 42 ----
.../execution_time/test_request_handlers.py | 77 -------
.../task_sdk/execution_time/test_supervisor.py | 25 +--
ts-sdk/src/generated/supervisor.ts | 48 +---
35 files changed, 574 insertions(+), 1188 deletions(-)
diff --git
a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
index f908a32808c..be03a4e489b 100644
---
a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
+++
b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
@@ -23,7 +23,7 @@ import contextlib
import itertools
import json
from collections import defaultdict
-from collections.abc import Callable, Iterator, Sequence
+from collections.abc import Callable, Iterable, Iterator, Sequence
from typing import Annotated, Any, NoReturn, cast
from uuid import UUID
@@ -39,7 +39,7 @@ from sqlalchemy import and_, func, or_, tuple_, update
from sqlalchemy.engine import CursorResult
from sqlalchemy.exc import DataError, NoResultFound, SQLAlchemyError
from sqlalchemy.orm import contains_eager, joinedload
-from sqlalchemy.sql import Select, select
+from sqlalchemy.sql import select
from sqlalchemy.sql.dml import Update
from structlog.contextvars import bind_contextvars
@@ -95,10 +95,15 @@ from airflow.models.asset import AssetActive
from airflow.models.base import ID_LEN
from airflow.models.dag import DagModel
from airflow.models.dagrun import DagRun as DR
-from airflow.models.dynamic_region import SENTINEL_REGION_ID,
AmbiguousProducerError
+from airflow.models.dynamic_region import AmbiguousProducerError
from airflow.models.hitl import HITLDetail
from airflow.models.log import Log
-from airflow.models.task_coordinates import TaskCoordinateResolver,
public_map_index_expression
+from airflow.models.task_coordinates import (
+ TaskCoordinateResolver,
+ build_coordinate_filters,
+ get_public_region,
+ public_map_index_expression,
+)
from airflow.models.taskinstance import TaskInstance as TI,
_stop_remaining_tasks
from airflow.models.taskreschedule import TaskReschedule
from airflow.models.trigger import Trigger, handle_event_submit
@@ -922,20 +927,37 @@ def ti_skip_downstream(
log.debug("Prepared task IDs for skipping", tasks=tasks)
resolver = TaskCoordinateResolver(dag_bag, session)
+ resolver.prefetch_regional_tasks(
+ dag_id, run_id, {task if isinstance(task, str) else task[0] for task
in tasks}
+ )
try:
# A bare task_id skips every TI of that task, so an already expanded
mapped task
# (e.g. one mapped over a literal list) is skipped too, not only
map_index -1.
+ targets = [(task, None) if isinstance(task, str) else (task[0],
task[1]) for task in tasks]
+ plain = [
+ (task_id, index)
+ for task_id, index in targets
+ if not resolver.has_regions(dag_id, run_id, task_id)
+ ]
selected_ids = {
ti.id
- for task in tasks
+ for task_id, index in targets
+ if (task_id, index) not in plain
for ti in resolver.resolve(
- dag_id=dag_id,
- run_id=run_id,
- task_id=task if isinstance(task, str) else task[0],
- caller=caller,
- map_indexes=None if isinstance(task, str) else task[1],
+ dag_id=dag_id, run_id=run_id, task_id=task_id, caller=caller,
map_indexes=index
)
}
+ if plain:
+ selected_ids.update(
+ session.scalars(
+ resolver.select_legacy_task_ids(
+ dag_id=dag_id,
+ run_id=run_id,
+ task_ids=[task_id for task_id, index in plain if index
is None],
+ slots=[(task_id, index) for task_id, index in plain if
index is not None],
+ )
+ )
+ )
except AmbiguousProducerError as error:
raise HTTPException(status.HTTP_409_CONFLICT, str(error)) from error
except ValueError as error:
@@ -1280,21 +1302,43 @@ async def get_previous_successful_dagrun(
return PrevSuccessfulDagRunResponse.model_validate(dag_run)
-def _filter_task_coordinates(
- query: Select, *, map_index: int | None, region_id: UUID | None,
region_index: int | None
-) -> Select:
- if region_index is not None and region_id is None:
- raise HTTPException(status.HTTP_400_BAD_REQUEST, "region_index
requires region_id")
- if region_id is not None:
- query = query.where(TI.region_id == region_id)
- if region_index is not None:
- query = query.where(TI.region_index == region_index)
- if map_index is not None:
- query = query.where(public_map_index_expression(TI) == map_index)
- return query
+def _find_superseded_ids(rows: Iterable[tuple[TI, int]], resolver:
TaskCoordinateResolver) -> set[UUID]:
+ """
+ Find live rows that a later loop pass of the same task slot supersedes.
+
+ A task inside a loop keeps every pass live, so callers that name a task
without a pass get its latest
+ one. Two live rows that are not told apart by their loop pass leave the
slot ambiguous.
+ """
+ slots: dict[tuple[str, str, int], dict[UUID, TI]] = defaultdict(dict)
+ for ti, public_index in rows:
+ slots[ti.run_id, ti.task_id, public_index][ti.id] = ti
+ superseded: set[UUID] = set()
+ for live in slots.values():
+ if len(live) < 2:
+ continue
+ try:
+ passes = {
+ ti_id: resolver.get_loop_pass(ti,
dag_version_id=ti.dag_version_id)
+ for ti_id, ti in live.items()
+ }
+ except ValueError as error:
+ raise HTTPException(
+ status.HTTP_409_CONFLICT, "Multiple live task instances share
a task slot"
+ ) from error
+ latest = max(passes.values(), key=lambda iteration: -1 if iteration is
None else iteration)
+ if sum(iteration == latest for iteration in passes.values()) > 1:
+ raise HTTPException(status.HTTP_409_CONFLICT, "Multiple live task
instances share a task slot")
+ superseded.update(ti_id for ti_id, iteration in passes.items() if
iteration != latest)
+ return superseded
[email protected]("/count", status_code=status.HTTP_200_OK)
[email protected](
+ "/count",
+ status_code=status.HTTP_200_OK,
+ responses=create_openapi_http_exception_doc(
+ [(status.HTTP_409_CONFLICT, "Multiple live task instances share a task
slot")]
+ ),
+)
def get_task_instance_count(
dag_id: str,
session: SessionDep,
@@ -1305,24 +1349,18 @@ def get_task_instance_count(
logical_dates: Annotated[list[UtcDateTime] | None, Query()] = None,
run_ids: Annotated[list[str] | None, Query()] = None,
states: Annotated[list[str] | None, Query()] = None,
- region_id: UUID | None = None,
- region_index: int | None = None,
) -> int:
- """Get the count of task instances matching the given criteria."""
- query = select(func.count()).select_from(TI).where(TI.dag_id == dag_id)
+ """Get the count of task instances matching the given criteria, counting a
looped task's latest pass."""
+ conditions = [TI.dag_id == dag_id, *build_coordinate_filters(TI,
map_index=map_index)]
if task_ids:
- query = query.where(TI.task_id.in_(task_ids))
-
- query = _filter_task_coordinates(
- query, map_index=map_index, region_id=region_id,
region_index=region_index
- )
+ conditions.append(TI.task_id.in_(task_ids))
if logical_dates:
- query = query.where(TI.logical_date.in_(logical_dates))
+ conditions.append(TI.logical_date.in_(logical_dates))
if run_ids:
- query = query.where(TI.run_id.in_(run_ids))
+ conditions.append(TI.run_id.in_(run_ids))
if task_group_id:
group_tasks = _get_group_tasks(
@@ -1332,31 +1370,44 @@ def get_task_instance_count(
if not group_tasks:
# If no task group tasks found, default to checking the task group
ID itself
# This matches the behavior in _get_external_task_group_task_ids
- query = query.where(TI.task_id == task_group_id,
public_map_index_expression(TI) == -1)
+ conditions.extend([TI.task_id == task_group_id,
public_map_index_expression(TI) == -1])
else:
- query = query.where(TI.id.in_(ti.id for ti in group_tasks))
+ conditions.append(TI.id.in_(ti.id for ti in group_tasks))
+
+ slot = public_map_index_expression(TI)
+ slots = (
+ select(TI.run_id.label("run_id"), TI.task_id.label("task_id"),
slot.label("slot"))
+ .where(*conditions)
+ .subquery()
+ )
+ shared_slots = session.execute(
+ select(slots.c.run_id, slots.c.task_id, slots.c.slot)
+ .group_by(slots.c.run_id, slots.c.task_id, slots.c.slot)
+ .having(func.count() > 1)
+ ).all()
+ if shared_slots:
+ rows = session.execute(
+ select(TI, slot).where(
+ *conditions, tuple_(TI.run_id, TI.task_id,
slot).in_([tuple(slot_) for slot_ in shared_slots])
+ )
+ ).tuples()
+ if superseded := _find_superseded_ids(rows,
TaskCoordinateResolver(dag_bag, session)):
+ conditions.append(TI.id.not_in(superseded))
if states:
if "null" in states:
not_none_states = [s for s in states if s != "null"]
if not_none_states:
- query = query.where(or_(TI.state.is_(None),
TI.state.in_(not_none_states)))
+ conditions.append(or_(TI.state.is_(None),
TI.state.in_(not_none_states)))
else:
- query = query.where(TI.state.is_(None))
+ conditions.append(TI.state.is_(None))
else:
- query = query.where(TI.state.in_(states))
+ conditions.append(TI.state.in_(states))
- count = session.scalar(query)
- return count or 0
+ return
session.scalar(select(func.count()).select_from(TI).where(*conditions)) or 0
[email protected](
- "/previous/{dag_id}/{task_id}",
- status_code=status.HTTP_200_OK,
- responses=create_openapi_http_exception_doc(
- [(status.HTTP_409_CONFLICT, "Explicit region coordinates are required
to select the previous task")]
- ),
-)
[email protected]("/previous/{dag_id}/{task_id}", status_code=status.HTTP_200_OK)
async def get_previous_task_instance(
dag_id: str,
task_id: str,
@@ -1364,11 +1415,9 @@ async def get_previous_task_instance(
logical_date: Annotated[UtcDateTime | None, Query()] = None,
map_index: Annotated[int, Query()] = -1,
state: Annotated[TaskInstanceState | None, Query()] = None,
- region_id: UUID | None = None,
- region_index: int | None = None,
) -> PreviousTIResponse | None:
"""
- Get the previous task instance matching the given criteria.
+ Get the previous task instance matching the given criteria, preferring a
looped task's latest pass.
:param dag_id: DAG ID (from path)
:param task_id: Task ID (from path)
@@ -1377,18 +1426,12 @@ async def get_previous_task_instance(
:param state: If provided, filters by TaskInstance state
"""
query = (
- select(TI, public_map_index_expression(TI),
func.count().over(partition_by=TI.run_id))
+ select(TI, public_map_index_expression(TI))
.where(TI.working_set.is_(True))
.join(DR, (TI.dag_id == DR.dag_id) & (TI.run_id == DR.run_id))
.options(contains_eager(TI.dag_run).load_only(DR.logical_date))
- .where(TI.dag_id == dag_id, TI.task_id == task_id)
- .order_by(DR.logical_date.desc())
- )
- query = _filter_task_coordinates(
- query,
- map_index=None if region_id is not None and region_index is not None
else map_index,
- region_id=region_id,
- region_index=region_index,
+ .where(TI.dag_id == dag_id, TI.task_id == task_id,
*build_coordinate_filters(TI, map_index=map_index))
+ .order_by(DR.logical_date.desc(), TI.region_index.desc())
)
if logical_date:
@@ -1401,11 +1444,8 @@ async def get_previous_task_instance(
row = (await session.execute(query.limit(1))).first()
if row is None:
return None
- ti, public_index, run_count = row
- if run_count > 1:
- raise HTTPException(
- status.HTTP_409_CONFLICT, "Select explicit region coordinates for
the previous task"
- )
+ ti, public_index = row
+ region_id, region_index = get_public_region(ti.region_id, ti.region_index)
return PreviousTIResponse(
task_id=ti.task_id,
@@ -1417,8 +1457,8 @@ async def get_previous_task_instance(
state=ti.state,
try_number=ti.try_number,
map_index=public_index,
- region_id=ti.region_id,
- region_index=ti.region_index,
+ region_id=region_id,
+ region_index=region_index,
duration=ti.duration,
)
@@ -1427,7 +1467,7 @@ async def get_previous_task_instance(
"/states",
status_code=status.HTTP_200_OK,
responses=create_openapi_http_exception_doc(
- [(status.HTTP_409_CONFLICT, "Multiple live task instances share a task
coordinate")]
+ [(status.HTTP_409_CONFLICT, "Multiple live task instances share a task
slot")]
),
)
def get_task_instance_states(
@@ -1439,13 +1479,12 @@ def get_task_instance_states(
task_group_id: Annotated[str | None, Query()] = None,
logical_dates: Annotated[list[UtcDateTime] | None, Query()] = None,
run_ids: Annotated[list[str] | None, Query()] = None,
- region_id: UUID | None = None,
- region_index: int | None = None,
) -> TaskStatesResponse:
- """Get the states for Task Instances with the given criteria."""
+ """Get the states for Task Instances with the given criteria, reporting a
looped task's latest pass."""
run_id_task_state_map: dict[str, dict[str, Any]] = defaultdict(dict)
- query = select(TI, public_map_index_expression(TI)).where(TI.dag_id ==
dag_id)
+ coordinates = build_coordinate_filters(TI, map_index=map_index)
+ query = select(TI, public_map_index_expression(TI)).where(TI.dag_id ==
dag_id, *coordinates)
if task_ids:
query = query.where(TI.task_id.in_(task_ids))
@@ -1456,32 +1495,24 @@ def get_task_instance_states(
if run_ids:
query = query.where(TI.run_id.in_(run_ids))
- query = _filter_task_coordinates(
- query, map_index=map_index, region_id=region_id,
region_index=region_index
- )
- results = session.execute(query).all()
+ results = session.execute(query).tuples().all()
if task_group_id:
group_tasks = _get_group_tasks(
dag_id, task_group_id, session, dag_bag, logical_dates, run_ids,
map_index
)
- group_query = _filter_task_coordinates(
- select(TI, public_map_index_expression(TI)).where(TI.id.in_(ti.id
for ti in group_tasks)),
- map_index=map_index,
- region_id=region_id,
- region_index=region_index,
+ group_query = select(TI, public_map_index_expression(TI)).where(
+ TI.id.in_(ti.id for ti in group_tasks), *coordinates
)
- group_results = session.execute(group_query).all()
+ group_results = session.execute(group_query).tuples().all()
results = [*results, *group_results] if task_ids else group_results
- identities: dict[tuple[str, str], UUID] = {}
+ superseded = _find_superseded_ids(results, TaskCoordinateResolver(dag_bag,
session))
for task, public_index in results:
+ if task.id in superseded:
+ continue
key = task.task_id if public_index < 0 else
f"{task.task_id}_{public_index}"
- identity = task.run_id, key
- if identity in identities and identities[identity] != task.id:
- raise HTTPException(status.HTTP_409_CONFLICT, "Task states require
explicit region coordinates")
- identities[identity] = task.id
run_id_task_state_map[task.run_id][key] = task.state
return TaskStatesResponse(task_states=run_id_task_state_map)
@@ -1492,8 +1523,6 @@ async def get_task_instance_breadcrumbs(
dag_id: str,
run_id: str,
session: AsyncSessionDep,
- region_id: UUID | None = None,
- region_index: int | None = None,
) -> TaskBreadcrumbsResponse:
query = (
select(
@@ -1509,19 +1538,15 @@ async def get_task_instance_breadcrumbs(
.where(TI.dag_id == dag_id, TI.run_id == run_id,
TI.state.in_(TerminalTIState))
.order_by(TI.task_id, public_map_index_expression(TI), TI.region_id,
TI.region_index)
)
- result = (
- await session.execute(
- _filter_task_coordinates(query, map_index=None,
region_id=region_id, region_index=region_index)
- )
- ).mappings()
+ result = (await session.execute(query)).mappings()
def _iter_breadcrumbs() -> Iterator[dict[str, Any]]:
for row in result:
- yield {
- str(k): v
- for k, v in row.items()
- if row.region_id != SENTINEL_REGION_ID or k not in
{"region_id", "region_index"}
- }
+ breadcrumb = {str(k): v for k, v in row.items() if k not in
{"region_id", "region_index"}}
+ region_id, region_index = get_public_region(row.region_id,
row.region_index)
+ if region_id is not None:
+ breadcrumb.update(region_id=region_id,
region_index=region_index)
+ yield breadcrumb
return TaskBreadcrumbsResponse(breadcrumbs=_iter_breadcrumbs())
diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py
b/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py
index 4224f418a5c..7cc976c0553 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py
@@ -44,8 +44,8 @@ from airflow.models.dagrun import DagRun
from airflow.models.dynamic_region import SENTINEL_REGION_ID,
AmbiguousProducerError
from airflow.models.task_coordinates import (
TaskCoordinateResolver,
+ build_coordinate_filters,
enclosing_loop,
- public_map_index_expression,
)
from airflow.models.taskinstance import TaskInstance
from airflow.models.xcom import XCOM_RETURN_KEY, XComModel, xcom_entity
@@ -146,8 +146,6 @@ def _build_xcom_read(
dag_bag: DBDagBag,
token: TIToken,
map_index: int | None = None,
- region_id: UUID | None = None,
- region_index: int | None = None,
include_prior_dates: bool = False,
) -> Select:
"""Select the XCom rows of the producers visible to the calling task
instance."""
@@ -161,20 +159,10 @@ def _build_xcom_read(
include_prior_dates=include_prior_dates,
)
try:
- if region_index is not None and region_id is None:
- raise ValueError("region_index requires region_id")
- if region_id is not None and region_index is not None:
- return read(region_id=region_id, map_indexes=region_index)
- if region_id == SENTINEL_REGION_ID or not resolver.has_regions(
- dag_id, None if include_prior_dates else run_id, task_id
- ):
+ if not resolver.has_regions(dag_id, None if include_prior_dates else
run_id, task_id):
return read(region_id=SENTINEL_REGION_ID, map_indexes=map_index)
caller = session.get(TaskInstance, token.id)
if include_prior_dates:
- if region_id is not None:
- raise ValueError(
- "Prior-run lookup requires producer coordinates resolved
separately for each run"
- )
task = None
if session.scalar(select(DagRun.id).where(DagRun.dag_id == dag_id,
DagRun.run_id == run_id)):
with suppress(TaskNotFound):
@@ -187,9 +175,7 @@ def _build_xcom_read(
else None,
)
if task is not None and enclosing_loop(task) is not None:
- raise ValueError(
- "Prior-date loop lookup requires an explicit target run
and producer coordinates"
- )
+ raise ValueError("Prior-date lookups are not supported for
tasks inside a loop")
return read(region_id=SENTINEL_REGION_ID,
include_node_regions=True, map_indexes=map_index)
producers = resolver.resolve(
dag_id=dag_id,
@@ -197,7 +183,6 @@ def _build_xcom_read(
task_id=task_id,
caller=caller,
map_indexes=map_index,
- region_id=region_id,
)
return XComModel.get_many(
run_id=run_id,
@@ -218,8 +203,6 @@ def xcom_query(
session: SessionDep,
dag_bag: DagBagDep,
map_index: Annotated[int | None, Query()] = None,
- region_id: UUID | None = None,
- region_index: int | None = None,
token=CurrentTIToken,
) -> Select:
return _build_xcom_read(
@@ -231,8 +214,6 @@ def xcom_query(
dag_bag=dag_bag,
token=token,
map_index=map_index,
- region_id=region_id,
- region_index=region_index,
)
@@ -248,8 +229,6 @@ def get_mapped_xcom_by_index(
offset: int,
session: SessionDep,
dag_bag: DagBagDep,
- region_id: UUID | None = None,
- region_index: int | None = None,
token=CurrentTIToken,
) -> XComSequenceIndexResponse:
xcom_read = _build_xcom_read(
@@ -260,8 +239,6 @@ def get_mapped_xcom_by_index(
session=session,
dag_bag=dag_bag,
token=token,
- region_id=region_id,
- region_index=region_index,
)
entity = xcom_entity(xcom_read)
xcom_query = xcom_read
@@ -290,8 +267,6 @@ class GetXComSliceFilterParams(BaseModel):
stop: int | None = None
step: int | None = None
include_prior_dates: bool = False
- region_id: UUID | None = None
- region_index: int | None = None
def _get_sliced_query_or_empty(query: Select, low: int, high: int) -> Select:
@@ -323,8 +298,6 @@ def get_mapped_xcom_by_slice(
dag_bag=dag_bag,
token=token,
include_prior_dates=params.include_prior_dates,
- region_id=params.region_id,
- region_index=params.region_index,
)
entity = xcom_entity(xcom_read)
query = xcom_read
@@ -434,8 +407,6 @@ class GetXcomFilterParams(BaseModel):
map_index: int = -1
include_prior_dates: bool = False
offset: int | None = None
- region_id: UUID | None = None
- region_index: int | None = None
@router.get(
@@ -460,9 +431,7 @@ def get_xcom(
TaskInstance.dag_id == dag_id,
TaskInstance.run_id == run_id,
TaskInstance.task_id == task_id,
- *_coordinate_filters(
- map_index=params.map_index, region_id=params.region_id,
region_index=params.region_index
- ),
+ *build_coordinate_filters(TaskInstance,
map_index=params.map_index),
)
.limit(1)
.execution_options(include_all_attempts=True)
@@ -494,8 +463,6 @@ def get_xcom(
dag_bag=dag_bag,
token=token,
map_index=params.map_index if params.offset is None else None,
- region_id=params.region_id,
- region_index=params.region_index,
include_prior_dates=params.include_prior_dates,
)
entity = xcom_entity(xcom_read)
@@ -573,8 +540,6 @@ def set_xcom(
),
] = None,
map_index: Annotated[int, Query()] = -1,
- region_id: UUID | None = None,
- region_index: int | None = None,
dag_result: Annotated[bool, Query(description="Whether this XCom is a dag
result")] = False,
mapped_length: Annotated[
int | None, Query(ge=0, description="Number of mapped tasks this value
expands into")
@@ -630,8 +595,6 @@ def set_xcom(
run_id=run_id,
task_id=task_id,
map_index=map_index,
- region_id=region_id,
- region_index=region_index,
session=session,
),
serialize=False,
@@ -665,8 +628,6 @@ def delete_xcom(
task_id: str,
key: Annotated[str, Path(min_length=1)],
map_index: Annotated[int, Query()] = -1,
- region_id: UUID | None = None,
- region_index: int | None = None,
token=CurrentTIToken,
):
"""Delete a single XCom Value."""
@@ -676,8 +637,6 @@ def delete_xcom(
run_id=run_id,
task_id=task_id,
map_index=map_index,
- region_id=region_id,
- region_index=region_index,
session=session,
)
XComModel.delete_for_attempts(
@@ -704,7 +663,9 @@ def _find_writer_id(
TaskInstance.dag_id == dag_id,
TaskInstance.run_id == run_id,
TaskInstance.task_id == task_id,
- *_coordinate_filters(map_index=map_index, region_id=region_id,
region_index=region_index),
+ *build_coordinate_filters(
+ TaskInstance, map_index=map_index, region_id=region_id,
region_index=region_index
+ ),
)
own = session.scalar(
select(TaskInstance.id)
@@ -718,7 +679,7 @@ def _find_writer_id(
raise HTTPException(
status.HTTP_409_CONFLICT,
f"More than one live task instance matches {dag_id}.{task_id} in
run {run_id}; "
- "address one by its region coordinates",
+ "a task inside a loop can only write XComs of its own attempt",
)
return live[0] if live else None
@@ -750,15 +711,3 @@ def _get_writer_id(
detail={"reason": "not_found", "message": f"Task Instance
{task_id!r} not found in {run_id!r}"},
)
return owner
-
-
-def _coordinate_filters(*, map_index: int, region_id: UUID | None,
region_index: int | None) -> list:
- """Match a task instance by explicit region coordinates, or by its public
map index when none are given."""
- if region_index is not None and region_id is None:
- raise HTTPException(status.HTTP_400_BAD_REQUEST, "region_index
requires region_id")
- if region_id is None:
- return [public_map_index_expression(TaskInstance) == map_index]
- filters = [TaskInstance.region_id == region_id]
- if region_index is not None:
- filters.append(TaskInstance.region_index == region_index)
- return filters
diff --git
a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
index 2ef78bc7bdd..b87c99b7c27 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
@@ -62,9 +62,6 @@ from airflow.api_fastapi.execution_api.versions.v2026_10_30
import (
AddTerminalStateRetryReasonField,
IdentifyArchivedTaskStateUpdates,
)
-from airflow.api_fastapi.execution_api.versions.v2026_10_30_xcom_params import
(
- AddRegionSelectorsToXComFilterParams,
-)
bundle = VersionBundle(
HeadVersion(),
@@ -77,7 +74,6 @@ bundle = VersionBundle(
AddMultiTeamToTIRunContext,
AddStoppedTaskReport,
AddTaskInstanceRegionCoordinates,
- AddRegionSelectorsToXComFilterParams,
IdentifyArchivedTaskStateUpdates,
),
Version(
diff --git
a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30_xcom_params.py
b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30_xcom_params.py
deleted file mode 100644
index 345e12aded9..00000000000
---
a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30_xcom_params.py
+++ /dev/null
@@ -1,37 +0,0 @@
-# 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.
-
-from __future__ import annotations
-
-from cadwyn import VersionChange, schema
-
-# These models live next to the routes that use them, and the routes import
this version package, so the
-# changes that touch them cannot sit in ``v2026_10_30`` without a circular
import.
-from airflow.api_fastapi.execution_api.routes.xcoms import
GetXcomFilterParams, GetXComSliceFilterParams
-
-
-class AddRegionSelectorsToXComFilterParams(VersionChange):
- """Add the `region_id` and `region_index` fields to
GetXComSliceFilterParams and GetXcomFilterParams."""
-
- description = __doc__
-
- instructions_to_migrate_to_previous_version = (
- schema(GetXComSliceFilterParams).field("region_id").didnt_exist,
- schema(GetXComSliceFilterParams).field("region_index").didnt_exist,
- schema(GetXcomFilterParams).field("region_id").didnt_exist,
- schema(GetXcomFilterParams).field("region_index").didnt_exist,
- )
diff --git a/airflow-core/src/airflow/executors/workloads/task.py
b/airflow-core/src/airflow/executors/workloads/task.py
index 139882c020a..3604f9d2b54 100644
--- a/airflow-core/src/airflow/executors/workloads/task.py
+++ b/airflow-core/src/airflow/executors/workloads/task.py
@@ -52,6 +52,7 @@ class TaskInstanceDTO(TaskInstance):
@property
def key(self) -> TaskInstanceKey:
+ """Identify the attempt the way the ORM does, by its stored index;
``map_index`` is the public one."""
from airflow.models.taskinstancekey import TaskInstanceKey
return TaskInstanceKey(
@@ -59,7 +60,7 @@ class TaskInstanceDTO(TaskInstance):
task_id=self.task_id,
run_id=self.run_id,
try_number=self.try_number,
- map_index=self.map_index,
+ map_index=self.region_index if self.region_index is not None else
self.map_index,
)
diff --git a/airflow-core/src/airflow/models/task_coordinates.py
b/airflow-core/src/airflow/models/task_coordinates.py
index e41b44e0d62..abe8cea5e56 100644
--- a/airflow-core/src/airflow/models/task_coordinates.py
+++ b/airflow-core/src/airflow/models/task_coordinates.py
@@ -20,7 +20,7 @@ from typing import TYPE_CHECKING, Protocol
from uuid import UUID
import attrs
-from sqlalchemy import case, or_, select
+from sqlalchemy import case, or_, select, tuple_
from airflow.exceptions import TaskNotFound
from airflow.models.dagrun import DagRun
@@ -29,6 +29,8 @@ from airflow.models.dynamic_region import (
AmbiguousProducerError,
DynamicRegion,
ProducerContext,
+ load_region_ancestry,
+ loop_position,
resolve_current_producers,
)
from airflow.models.taskinstance import TaskInstance
@@ -37,6 +39,7 @@ from airflow.serialization.definitions.taskgroup import
SerializedLoopTaskGroup
if TYPE_CHECKING:
from collections.abc import Collection
+ from sqlalchemy import Select
from sqlalchemy.orm import Session
from sqlalchemy.sql.elements import ColumnElement
@@ -76,6 +79,27 @@ def public_map_index_expression(model) -> ColumnElement[int]:
)
+def get_public_region(region_id: UUID, region_index: int) -> tuple[UUID |
None, int | None]:
+ """Report the region a client sees: ``None`` for a task instance that
lives in no dynamic region."""
+ if region_id == SENTINEL_REGION_ID:
+ return None, None
+ return region_id, region_index
+
+
+def build_coordinate_filters(
+ model, *, map_index: int | None, region_id: UUID | None = None,
region_index: int | None = None
+) -> list[ColumnElement[bool]]:
+ """Match task instances by explicit region coordinates, or by their public
map index when none are given."""
+ if region_index is not None and region_id is None:
+ raise ValueError("region_index requires region_id")
+ if region_id is None:
+ return [] if map_index is None else
[public_map_index_expression(model) == map_index]
+ filters = [model.region_id == region_id]
+ if region_index is not None:
+ filters.append(model.region_index == region_index)
+ return filters
+
+
@attrs.define
class TaskCoordinateResolver:
"""Resolve task coordinates against the definitions pinned to their
executions."""
@@ -84,6 +108,7 @@ class TaskCoordinateResolver:
session: Session
_dags: dict[UUID, SerializedDAG] = attrs.field(factory=dict, init=False)
_region_nodes: dict[UUID, str | None] = attrs.field(factory=dict,
init=False)
+ _regional_tasks: dict[tuple[str, str | None, str], bool] =
attrs.field(factory=dict, init=False)
def get_task(
self, dag_id: str, run_id: str, task_id: str, *, dag_version_id: UUID
| None = None
@@ -131,7 +156,9 @@ class TaskCoordinateResolver:
for task in tasks
}
if len(classifications) != 1:
- raise AmbiguousProducerError("Producer definitions differ; select
explicit region coordinates")
+ raise AmbiguousProducerError(
+ "Producer definitions differ between the live task instances
of the requested task"
+ )
return tasks[0]
def public_map_index(self, ti: TaskCoordinate, *, dag_version_id: UUID |
None = None) -> int:
@@ -158,7 +185,69 @@ class TaskCoordinateResolver:
)
return self._region_nodes[ti.region_id] == ti.task_id
+ def prefetch_regional_tasks(self, dag_id: str, run_id: str, task_ids:
Collection[str]) -> None:
+ """Learn in one query which of ``task_ids`` have live task instances
in a dynamic region."""
+ if not task_ids:
+ return
+ regional = set(
+ self.session.scalars(
+ select(TaskInstance.task_id)
+ .where(
+ TaskInstance.working_set.is_(True),
+ TaskInstance.dag_id == dag_id,
+ TaskInstance.run_id == run_id,
+ TaskInstance.task_id.in_(task_ids),
+ TaskInstance.region_id != SENTINEL_REGION_ID,
+ )
+ .distinct()
+ )
+ )
+ for task_id in task_ids:
+ self._regional_tasks[dag_id, run_id, task_id] = task_id in regional
+
+ def select_legacy_task_ids(
+ self,
+ *,
+ dag_id: str,
+ run_id: str,
+ task_ids: Collection[str],
+ slots: Collection[tuple[str, int]],
+ ) -> Select[tuple[UUID]]:
+ """Select, in one statement, the region-less task instances that the
named tasks or slots address."""
+ targets = []
+ if task_ids:
+ targets.append(TaskInstance.task_id.in_(task_ids))
+ if slots:
+ targets.append(tuple_(TaskInstance.task_id,
TaskInstance.region_index).in_(list(slots)))
+ return select(TaskInstance.id).where(
+ TaskInstance.working_set.is_(True),
+ TaskInstance.dag_id == dag_id,
+ TaskInstance.run_id == run_id,
+ TaskInstance.region_id == SENTINEL_REGION_ID,
+ or_(*targets),
+ )
+
+ def get_loop_pass(self, ti: TaskCoordinate, *, dag_version_id: UUID | None
= None) -> int | None:
+ """Return the iteration of the loop enclosing ``ti``'s task, or
``None`` outside any loop."""
+ if ti.region_id == SENTINEL_REGION_ID:
+ return None
+ try:
+ loop = enclosing_loop(
+ self.get_task(ti.dag_id, ti.run_id, ti.task_id,
dag_version_id=dag_version_id)
+ )
+ except (TaskNotFound, ValueError):
+ return None
+ if loop is None or loop.group_id is None:
+ return None
+ regions = load_region_ancestry(
+ [ti.region_id], dag_id=ti.dag_id, run_id=ti.run_id,
session=self.session
+ )
+ position = loop_position(regions, ti.region_id, ti.region_index,
loop.group_id)
+ return None if position is None else position[1]
+
def has_regions(self, dag_id: str, run_id: str | None, task_id: str) ->
bool:
+ if (known := self._regional_tasks.get((dag_id, run_id, task_id))) is
not None:
+ return known
query = select(TaskInstance.id).where(
TaskInstance.working_set.is_(True),
TaskInstance.dag_id == dag_id,
diff --git a/airflow-core/src/airflow/utils/log/task_log_address.py
b/airflow-core/src/airflow/utils/log/task_log_address.py
index 91c94bebd2d..ff599def18c 100644
--- a/airflow-core/src/airflow/utils/log/task_log_address.py
+++ b/airflow-core/src/airflow/utils/log/task_log_address.py
@@ -61,13 +61,19 @@ def region_log_position(
*,
regions: Mapping[UUID, DynamicRegion],
node_kinds: Mapping[str, str],
+ task_id: str | None = None,
) -> str:
- """Render the loop nesting and fork position of one coordinate, empty
outside any loop."""
+ """
+ Render the loop nesting and fork position of one coordinate, empty outside
any loop.
+
+ A node the pinned definition no longer lists is classified from the stored
region: the task's own
+ region is its mapped expansion and any other region encloses it as a loop.
+ """
frames = []
inside_loop = False
while region_id != SENTINEL_REGION_ID:
region = regions[region_id]
- kind = node_kinds[region.node_id]
+ kind = node_kinds.get(region.node_id) or ("map" if region.node_id ==
task_id else "loop")
if kind not in {"loop", "map"}:
raise ValueError(f"Unsupported log-address region kind: {kind}")
position = 1
@@ -170,8 +176,7 @@ def prepare_task_log_contexts(
templates[None] =
session.scalar(select(LogTemplate.filename).order_by(LogTemplate.id).limit(1))
regional = [ti for ti in tis if ti.region_id != SENTINEL_REGION_ID]
regions: dict[UUID, DynamicRegion] = {}
- dags = {}
- node_kinds = {}
+ node_kinds: dict[UUID, dict[str, str]] = {}
if regional:
frontier = {ti.region_id for ti in regional}
while frontier:
@@ -191,7 +196,6 @@ def prepare_task_log_contexts(
):
row.load_op_links = False
dag = row.dag
- dags[row.dag_version_id] = dag
node_kinds[row.dag_version_id] = {
group_id: "loop"
for group_id, group in
dag.task_group.get_task_group_dict().items()
@@ -210,13 +214,14 @@ def prepare_task_log_contexts(
map_index = ti.map_index
if ti.region_id != SENTINEL_REGION_ID:
version_id = ti.dag_version_id or run.created_dag_version_id
- if version_id is None:
- raise ValueError(f"Regional task instance {ti.id} has no
pinned DAG version")
- dag = dags[version_id]
position = region_log_position(
- ti.region_id, ti.region_index, regions=regions,
node_kinds=node_kinds[version_id]
+ ti.region_id,
+ ti.region_index,
+ regions=regions,
+ node_kinds=node_kinds.get(version_id, {}) if version_id else
{},
+ task_id=ti.task_id,
)
- map_index = ti.region_index if
dag.get_task(ti.task_id).get_needs_expansion() else -1
+ map_index = ti.region_index if regions[ti.region_id].node_id ==
ti.task_id else -1
contexts[ti.id] = TaskLogContext(
template,
(run.logical_date or run.run_after).isoformat(),
diff --git
a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
index 6c7641c5ac7..d9a4bbdb525 100644
---
a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
+++
b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
@@ -61,7 +61,6 @@ from airflow.models.task_state_store import
TaskStateStoreModel
from airflow.models.taskinstance import TaskInstance, clear_task_instances
from airflow.models.xcom import XComModel
from airflow.providers.standard.operators.empty import EmptyOperator
-from airflow.providers.standard.operators.python import PythonOperator
from airflow.sdk import Asset, TaskGroup, TriggerRule, task, task_group
from airflow.sdk.definitions._internal.loop import create_loop
from airflow.state.metastore import MetastoreBackend
@@ -195,67 +194,96 @@ def loop_reader_tis(dag_maker, session):
return current, previous
-def test_execution_count_uses_public_index_and_exact_region_scope(client,
loop_reader_tis):
+def test_execution_count_reports_the_latest_loop_pass(client, loop_reader_tis):
current, _ = loop_reader_tis
params = {"dag_id": current.dag_id, "task_ids": [current.task_id],
"map_index": -1}
url = "/execution/task-instances/count"
- assert client.get(url, params=params).json() == 2
- assert (
- client.get(url, params={**params, "region_id": str(current.region_id),
"region_index": 2}).json() == 1
- )
+ assert client.get(url, params=params).json() == 1
+ assert client.get(url, params={**params, "states": ["success"]}).json() ==
1
+ assert client.get(url, params={**params, "states": ["failed"]}).json() == 0
+
+
[email protected]("older_row_is_last", [False, True])
+def test_execution_states_report_the_latest_loop_pass(client, session,
loop_reader_tis, older_row_is_last):
+ current, previous = loop_reader_tis
+ latest = current
+ if older_row_is_last:
+ current.region_index, previous.region_index = 0, 2
+ latest = previous
+ session.commit()
+ params = {"dag_id": current.dag_id, "task_ids": [current.task_id],
"map_index": -1}
+ response = client.get("/execution/task-instances/states", params=params)
[email protected]("endpoint", ["states", "previous"])
-def
test_execution_reader_rejects_loop_collision_and_accepts_exact_scope(client,
loop_reader_tis, endpoint):
- current, _ = loop_reader_tis
- params = {"dag_id": current.dag_id, "task_ids": [current.task_id],
"map_index": -1}
- url = "/execution/task-instances/states"
- if endpoint == "previous":
- url =
f"/execution/task-instances/previous/{current.dag_id}/{current.task_id}"
- params = {"map_index": -1}
- assert client.get(url, params=params).status_code == 409
- response = client.get(url, params={**params, "region_id":
str(current.region_id), "region_index": 2})
assert response.status_code == 200
- if endpoint == "states":
- assert response.json() == {"task_states": {current.run_id:
{current.task_id: "success"}}}
- else:
- assert response.json()["map_index"] == -1
- assert response.json()["region_index"] == 2
+ assert response.json() == {"task_states": {current.run_id:
{current.task_id: latest.state}}}
-def
test_execution_breadcrumbs_keep_regional_identity_separate_from_map_index(client,
loop_reader_tis):
[email protected]("endpoint", ["count", "states"])
+def
test_execution_reads_reject_live_task_instances_in_one_slot_that_the_loop_cannot_tell_apart(
+ client, session, loop_reader_tis, endpoint
+):
current, _ = loop_reader_tis
+ stray = TaskInstance(task=current.task, run_id=current.run_id,
dag_version_id=current.dag_version_id)
+ stray.region_id, stray.region_index, stray.state = uuid4(), 5,
State.SUCCESS
+ session.add(stray)
+ session.commit()
+
response = client.get(
- "/execution/task-instances/breadcrumbs", params={"dag_id":
current.dag_id, "run_id": current.run_id}
+ f"/execution/task-instances/{endpoint}",
+ params={"dag_id": current.dag_id, "task_ids": [current.task_id],
"map_index": -1},
)
- assert response.status_code == 200
- breadcrumbs = response.json()["breadcrumbs"]
- assert len(breadcrumbs) == 2
- assert [row["map_index"] for row in breadcrumbs] == [-1, -1]
- assert {row["region_index"] for row in breadcrumbs} == {1, 2}
+ assert response.status_code == 409
-def test_previous_ti_full_coordinates_override_default_public_index(client,
dag_maker, session):
- with dag_maker(serialized=True):
- PythonOperator.partial(task_id="mapped",
python_callable=str).expand(op_args=[[0], [1], [2]])
- dr = dag_maker.create_dagrun()
- region = DynamicRegion(dag_id=dr.dag_id, run_id=dr.run_id,
node_id="mapped")
- session.add(region)
+
[email protected]("endpoint", ["count", "states"])
+def test_execution_reads_reject_two_live_task_instances_in_one_loop_pass(
+ client, session, loop_reader_tis, endpoint
+):
+ current, _ = loop_reader_tis
+ fork = DynamicRegion(
+ dag_id=current.dag_id, run_id=current.run_id, node_id="body",
forked_from_region_id=current.region_id
+ )
+ session.add(fork)
session.flush()
- for ti in dr.task_instances:
- ti.region_id = region.id
+ twin = TaskInstance(task=current.task, run_id=current.run_id,
dag_version_id=current.dag_version_id)
+ twin.region_id, twin.region_index, twin.state = fork.id, 2, State.SUCCESS
+ session.add(twin)
session.commit()
response = client.get(
- f"/execution/task-instances/previous/{dr.dag_id}/mapped",
- params={"region_id": str(region.id), "region_index": 2},
+ f"/execution/task-instances/{endpoint}",
+ params={"dag_id": current.dag_id, "task_ids": [current.task_id],
"map_index": -1},
+ )
+
+ assert response.status_code == 409
+
+
+def test_execution_previous_reports_the_latest_loop_pass(client,
loop_reader_tis):
+ current, _ = loop_reader_tis
+
+ response = client.get(
+
f"/execution/task-instances/previous/{current.dag_id}/{current.task_id}",
params={"map_index": -1}
)
assert response.status_code == 200
- assert response.json()["map_index"] == 2
+ assert response.json()["map_index"] == -1
assert response.json()["region_index"] == 2
+def
test_execution_breadcrumbs_keep_regional_identity_separate_from_map_index(client,
loop_reader_tis):
+ current, _ = loop_reader_tis
+ response = client.get(
+ "/execution/task-instances/breadcrumbs", params={"dag_id":
current.dag_id, "run_id": current.run_id}
+ )
+ assert response.status_code == 200
+ breadcrumbs = response.json()["breadcrumbs"]
+ assert len(breadcrumbs) == 2
+ assert [row["map_index"] for row in breadcrumbs] == [-1, -1]
+ assert {row["region_index"] for row in breadcrumbs} == {1, 2}
+
+
class TestTIRunState:
@pytest.mark.parametrize("regional", [False, True])
def test_startup_xcom_cleanup_keys_belong_to_the_starting_try(
@@ -3316,6 +3344,49 @@ class TestTISkipDownstream:
def teardown_method(self):
clear_db_runs()
+ def test_ti_skip_downstream_asks_once_which_tasks_have_regions(self,
client, session, dag_maker):
+ with dag_maker("skip_downstream_many", session=session):
+ head = EmptyOperator(task_id="head")
+ for index in range(6):
+ head >> EmptyOperator(task_id=f"t{index}")
+ dr = dag_maker.create_dagrun(run_id="run")
+ head_ti = dr.get_task_instance("head")
+ head_ti.set_state(State.SUCCESS)
+ session.commit()
+
+ with capture_orm_selects("task_instance") as statements:
+ response = client.patch(
+ f"/execution/task-instances/{head_ti.id}/skip-downstream",
+ json={"tasks": [f"t{index}" for index in range(6)]},
+ )
+
+ assert response.status_code == 204
+ region_lookups = [sql for sql in statements if "region_id !=" in sql]
+ assert len(region_lookups) == 1
+
+ def test_ti_skip_downstream_with_hundreds_of_targets_uses_constant_queries(
+ self, client, session, dag_maker
+ ):
+ with dag_maker("skip_downstream_hundreds", session=session):
+ head = EmptyOperator(task_id="head")
+ for index in range(520):
+ head >> EmptyOperator(task_id=f"t{index}")
+ dr = dag_maker.create_dagrun(run_id="run")
+ head_ti = dr.get_task_instance("head")
+ head_ti.set_state(State.SUCCESS)
+ session.commit()
+
+ with capture_orm_selects("task_instance") as statements:
+ response = client.patch(
+ f"/execution/task-instances/{head_ti.id}/skip-downstream",
+ json={"tasks": [f"t{index}" for index in range(520)]},
+ )
+
+ assert response.status_code == 204
+ assert len(statements) <= 6
+ session.expire_all()
+ assert {ti.state for ti in dr.get_task_instances() if ti.task_id !=
"head"} == {State.SKIPPED}
+
@pytest.mark.parametrize("_json", (({"tasks": ["t1"]}), ({"tasks": [("t1",
-1)]})))
def test_ti_skip_downstream(self, client, session, create_task_instance,
dag_maker, _json):
with dag_maker("skip_downstream_dag", session=session):
@@ -4415,6 +4486,31 @@ class TestGetPreviousTI:
assert data["run_id"] == "run1"
assert data["state"] == State.SUCCESS
+ def test_get_previous_ti_reports_no_region_for_a_task_outside_any_region(
+ self, client, session, create_task_instance
+ ):
+ create_task_instance(
+ task_id="test_task",
+ state=State.SUCCESS,
+ logical_date=timezone.datetime(2025, 1, 1),
+ run_id="run1",
+ )
+ create_task_instance(
+ task_id="test_task",
+ state=State.SUCCESS,
+ logical_date=timezone.datetime(2025, 1, 2),
+ run_id="run2",
+ )
+ session.commit()
+
+ response = client.get(
+ "/execution/task-instances/previous/dag/test_task",
+ params={"logical_date": "2025-01-02T00:00:00Z"},
+ )
+
+ assert response.json()["region_id"] is None
+ assert response.json()["region_index"] is None
+
def test_get_previous_ti_with_state_filter(self, client, session,
create_task_instance):
"""Test get_previous_ti with state filter."""
# Create TIs with different states
diff --git
a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py
b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py
index 59d00a28c0b..a50fde183bb 100644
---
a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py
+++
b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py
@@ -93,27 +93,6 @@ def
test_regional_mapped_xcom_reads_resolve_live_producers_before_slicing(client
assert response.json() == expected[suffix]
-def test_xcom_write_coordinates_must_match_the_calling_task_instance(
- client, create_task_instance, authenticate_as, session
-):
- ti = create_task_instance()
- region = DynamicRegion(dag_id=ti.dag_id, run_id=ti.run_id, node_id="loop")
- session.add(region)
- session.commit()
- authenticate_as(ti)
- url = f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/key"
- other = {"region_id": str(region.id), "region_index": -1}
- own = {"region_id": str(ti.region_id), "region_index": ti.region_index}
-
- assert client.post(url, params=other, json="regional").status_code == 404
- assert client.delete(url, params=other).status_code == 404
- assert client.post(url, params=own, json="legacy").status_code == 201
- assert client.delete(url, params=other).status_code == 404
- assert client.get(url, params=own).json() == {"key": "key", "value":
"legacy"}
- assert client.get(url, params=other).status_code == 404
- assert client.delete(url, params=own).status_code == 200
-
-
@pytest.fixture
def loop_xcoms(dag_maker, session, authenticate_as):
@task_group
@@ -151,26 +130,21 @@ def loop_xcoms(dag_maker, session, authenticate_as):
def test_loop_xcom_omission_uses_consumer_pass_and_retained_region(client,
loop_xcoms):
- dr, producer, consumer, previous = loop_xcoms
+ dr, producer, consumer, _ = loop_xcoms
base = f"/execution/xcoms/{dr.dag_id}/{dr.run_id}"
assert client.get(f"{base}/{producer.task_id}/key").json() == {"key":
"key", "value": "current"}
- assert client.get(
- f"{base}/{producer.task_id}/key", params={"region_id":
str(previous.region_id), "region_index": 1}
- ).json() == {"key": "key", "value": "previous"}
assert client.post(f"{base}/{consumer.task_id}/key",
json="consumer").status_code == 201
- assert client.get(
- f"{base}/{consumer.task_id}/key", params={"region_id":
str(consumer.region_id), "region_index": 2}
- ).json() == {"key": "key", "value": "consumer"}
+ assert client.get(f"{base}/{consumer.task_id}/key").json() == {"key":
"key", "value": "consumer"}
-def
test_loop_prior_dates_requires_explicit_iteration_but_allows_outside_producer(client,
loop_xcoms):
+def
test_loop_prior_dates_are_unsupported_but_allowed_for_outside_producer(client,
loop_xcoms):
dr, producer, _, _ = loop_xcoms
base = f"/execution/xcoms/{dr.dag_id}/{dr.run_id}"
response = client.get(f"{base}/{producer.task_id}/key",
params={"include_prior_dates": True})
assert response.status_code == 400
- assert "explicit target run" in response.json()["detail"]
+ assert "not supported for tasks inside a loop" in response.json()["detail"]
assert client.get(f"{base}/outside/key", params={"include_prior_dates":
True}).json() == {
"key": "key",
"value": "outside",
diff --git a/airflow-core/tests/unit/executors/test_workloads.py
b/airflow-core/tests/unit/executors/test_workloads.py
index 0ef36d27b48..4dbc5e01962 100644
--- a/airflow-core/tests/unit/executors/test_workloads.py
+++ b/airflow-core/tests/unit/executors/test_workloads.py
@@ -391,6 +391,8 @@ def
test_execute_task_projects_unmapped_loop_index_without_changing_storage(dag_
assert workload.ti.map_index == -1
assert workload.ti.region_index == 2
assert workload.ti.region_id == region.id
+ assert workload.key == ti.key
+ assert workload.key.map_index == 2
assert "pass=2" in str(workload.log_path)
assert ti.region_index == 2
diff --git a/airflow-core/tests/unit/models/test_task_coordinates.py
b/airflow-core/tests/unit/models/test_task_coordinates.py
index 429bc59d4ba..c390430c55a 100644
--- a/airflow-core/tests/unit/models/test_task_coordinates.py
+++ b/airflow-core/tests/unit/models/test_task_coordinates.py
@@ -17,13 +17,14 @@
from __future__ import annotations
import pytest
+from sqlalchemy import select
from airflow.executors.workloads.base import BundleInfo
from airflow.executors.workloads.task import ExecuteTask
from airflow.models.dag_version import DagVersion
from airflow.models.dagbag import DBDagBag
from airflow.models.dynamic_region import AmbiguousProducerError, DynamicRegion
-from airflow.models.task_coordinates import TaskCoordinateResolver
+from airflow.models.task_coordinates import TaskCoordinateResolver,
build_coordinate_filters
from airflow.models.taskinstance import TaskInstance
from airflow.providers.standard.operators.empty import EmptyOperator
from airflow.providers.standard.operators.python import PythonOperator
@@ -120,6 +121,31 @@ def
test_loop_without_context_requires_explicit_scope(loop_coordinates, session)
) == (producer,)
+def
test_coordinate_filters_match_the_exact_region_or_the_public_index(loop_coordinates,
session):
+ _, producer, _, previous, _ = loop_coordinates
+ task = TaskInstance.task_id == producer.task_id
+
+ by_public_index = session.scalars(
+ select(TaskInstance.id).where(task,
*build_coordinate_filters(TaskInstance, map_index=-1))
+ ).all()
+ exact = session.scalars(
+ select(TaskInstance.id).where(
+ task,
+ *build_coordinate_filters(
+ TaskInstance, map_index=-1, region_id=producer.region_id,
region_index=producer.region_index
+ ),
+ )
+ ).all()
+
+ assert set(by_public_index) == {producer.id, previous.id}
+ assert exact == [producer.id]
+
+
+def test_coordinate_filters_require_a_region_id_with_a_region_index():
+ with pytest.raises(ValueError, match="region_index requires region_id"):
+ build_coordinate_filters(TaskInstance, map_index=-1, region_index=2)
+
+
def test_sentinel_consumer_resolves_modern_mapped_producers(dag_maker,
session):
with dag_maker(serialized=True):
PythonOperator.partial(task_id="mapped",
python_callable=str).expand(op_args=[[1], [2]])
diff --git a/airflow-core/tests/unit/utils/log/test_task_log_address.py
b/airflow-core/tests/unit/utils/log/test_task_log_address.py
index ea4df810eb9..156def635e5 100644
--- a/airflow-core/tests/unit/utils/log/test_task_log_address.py
+++ b/airflow-core/tests/unit/utils/log/test_task_log_address.py
@@ -23,6 +23,7 @@ import pytest
from sqlalchemy import event
from airflow.configuration import conf
+from airflow.models.dag_version import DagVersion
from airflow.models.dynamic_region import DynamicRegion
from airflow.models.taskinstance import TaskInstance
from airflow.providers.standard.operators.empty import EmptyOperator
@@ -242,3 +243,37 @@ def
test_mapped_log_context_projects_expansion_position(dag_maker, session, insi
assert path.endswith(f"map_index=0/{token}/attempt=1.log")
else:
assert path.endswith("task_id=body.mapped/map_index=0/attempt=1.log")
+
+
[email protected]_test
+def
test_log_context_falls_back_to_stored_regions_when_the_pinned_dag_lost_the_nodes(dag_maker,
session):
+ @task_group(group_id="body")
+ def body():
+ PythonOperator.partial(task_id="mapped",
python_callable=str).expand(op_args=[[1], [2]])
+
+ with dag_maker("loop_dag", serialized=True):
+ create_loop(body, max_iterations=3)
+ run = dag_maker.create_dagrun()
+ ti = next(ti for ti in run.task_instances if ti.task_id == "body.mapped")
+ loop_region = DynamicRegion(dag_id=run.dag_id, run_id=run.run_id,
node_id="body")
+ session.add(loop_region)
+ session.flush()
+ expansion = DynamicRegion(
+ dag_id=run.dag_id,
+ run_id=run.run_id,
+ node_id=ti.task_id,
+ parent_region_id=loop_region.id,
+ parent_region_index=2,
+ )
+ session.add(expansion)
+ with dag_maker("replacement_dag", serialized=True):
+ EmptyOperator(task_id="unrelated")
+ ti.dag_version_id = DagVersion.get_latest_version("replacement_dag",
session=session).id
+ ti.region_id, ti.region_index = expansion.id, 0
+ session.flush()
+ ti.task = None
+
+ context = prepare_task_log_contexts([ti], session=session)[ti.id]
+
+ assert context.log_position == "pass=2/map=0"
+ assert context.map_index == 0
diff --git a/go-sdk/airflow/spec_fields.gen.go
b/go-sdk/airflow/spec_fields.gen.go
index 2d2c25c4f9d..26e552143ed 100644
--- a/go-sdk/airflow/spec_fields.gen.go
+++ b/go-sdk/airflow/spec_fields.gen.go
@@ -75,7 +75,6 @@ var taskSpecFields = map[string]schemaField{
var taskGroupSpecFields = map[string]schemaField{
"DocMD": {key: "doc_md"},
"GroupDisplayName": {key: "group_display_name"},
- "Loop": {key: "loop"},
"PrefixGroupID": {key: "prefix_group_id"},
"Tooltip": {key: "tooltip"},
"UIColor": {key: "ui_color"},
diff --git a/go-sdk/pkg/execution/genmodels/models.gen.go
b/go-sdk/pkg/execution/genmodels/models.gen.go
index 62de3f6505e..a56ac896b9c 100644
--- a/go-sdk/pkg/execution/genmodels/models.gen.go
+++ b/go-sdk/pkg/execution/genmodels/models.gen.go
@@ -736,12 +736,6 @@ type DeleteXCom struct {
// MapIndex corresponds to the JSON schema field "map_index".
MapIndex interface{} `msgpack:"map_index,omitempty"`
- // RegionID corresponds to the JSON schema field "region_id".
- RegionID interface{} `msgpack:"region_id,omitempty"`
-
- // RegionIndex corresponds to the JSON schema field "region_index".
- RegionIndex interface{} `msgpack:"region_index,omitempty"`
-
// RunID corresponds to the JSON schema field "run_id".
RunID string `msgpack:"run_id"`
@@ -1021,12 +1015,6 @@ type GetPreviousTI struct {
// MapIndex corresponds to the JSON schema field "map_index".
MapIndex *int `msgpack:"map_index,omitempty"`
- // RegionID corresponds to the JSON schema field "region_id".
- RegionID interface{} `msgpack:"region_id,omitempty"`
-
- // RegionIndex corresponds to the JSON schema field "region_index".
- RegionIndex interface{} `msgpack:"region_index,omitempty"`
-
// State corresponds to the JSON schema field "state".
State interface{} `msgpack:"state,omitempty"`
@@ -1047,12 +1035,6 @@ type GetTICount struct {
// MapIndex corresponds to the JSON schema field "map_index".
MapIndex interface{} `msgpack:"map_index,omitempty"`
- // RegionID corresponds to the JSON schema field "region_id".
- RegionID interface{} `msgpack:"region_id,omitempty"`
-
- // RegionIndex corresponds to the JSON schema field "region_index".
- RegionIndex interface{} `msgpack:"region_index,omitempty"`
-
// RunIds corresponds to the JSON schema field "run_ids".
RunIds *RunIds `msgpack:"run_ids,omitempty"`
@@ -1073,12 +1055,6 @@ type GetTaskBreadcrumbs struct {
// DagID corresponds to the JSON schema field "dag_id".
DagID string `msgpack:"dag_id"`
- // RegionID corresponds to the JSON schema field "region_id".
- RegionID interface{} `msgpack:"region_id,omitempty"`
-
- // RegionIndex corresponds to the JSON schema field "region_index".
- RegionIndex interface{} `msgpack:"region_index,omitempty"`
-
// RunID corresponds to the JSON schema field "run_id".
RunID string `msgpack:"run_id"`
@@ -1118,12 +1094,6 @@ type GetTaskStates struct {
// MapIndex corresponds to the JSON schema field "map_index".
MapIndex interface{} `msgpack:"map_index,omitempty"`
- // RegionID corresponds to the JSON schema field "region_id".
- RegionID interface{} `msgpack:"region_id,omitempty"`
-
- // RegionIndex corresponds to the JSON schema field "region_index".
- RegionIndex interface{} `msgpack:"region_index,omitempty"`
-
// RunIds corresponds to the JSON schema field "run_ids".
RunIds *RunIds `msgpack:"run_ids,omitempty"`
@@ -1172,12 +1142,6 @@ type GetXCom struct {
// MapIndex corresponds to the JSON schema field "map_index".
MapIndex interface{} `msgpack:"map_index,omitempty"`
- // RegionID corresponds to the JSON schema field "region_id".
- RegionID interface{} `msgpack:"region_id,omitempty"`
-
- // RegionIndex corresponds to the JSON schema field "region_index".
- RegionIndex interface{} `msgpack:"region_index,omitempty"`
-
// RunID corresponds to the JSON schema field "run_id".
RunID string `msgpack:"run_id"`
@@ -1196,12 +1160,6 @@ type GetXComCount struct {
// Key corresponds to the JSON schema field "key".
Key string `msgpack:"key"`
- // RegionID corresponds to the JSON schema field "region_id".
- RegionID interface{} `msgpack:"region_id,omitempty"`
-
- // RegionIndex corresponds to the JSON schema field "region_index".
- RegionIndex interface{} `msgpack:"region_index,omitempty"`
-
// RunID corresponds to the JSON schema field "run_id".
RunID string `msgpack:"run_id"`
@@ -1222,12 +1180,6 @@ type GetXComSequenceItem struct {
// Offset corresponds to the JSON schema field "offset".
Offset int `msgpack:"offset"`
- // RegionID corresponds to the JSON schema field "region_id".
- RegionID interface{} `msgpack:"region_id,omitempty"`
-
- // RegionIndex corresponds to the JSON schema field "region_index".
- RegionIndex interface{} `msgpack:"region_index,omitempty"`
-
// RunID corresponds to the JSON schema field "run_id".
RunID string `msgpack:"run_id"`
@@ -1248,12 +1200,6 @@ type GetXComSequenceSlice struct {
// Key corresponds to the JSON schema field "key".
Key string `msgpack:"key"`
- // RegionID corresponds to the JSON schema field "region_id".
- RegionID interface{} `msgpack:"region_id,omitempty"`
-
- // RegionIndex corresponds to the JSON schema field "region_index".
- RegionIndex interface{} `msgpack:"region_index,omitempty"`
-
// RunID corresponds to the JSON schema field "run_id".
RunID string `msgpack:"run_id"`
@@ -1618,12 +1564,6 @@ type SetXCom struct {
// MappedLength corresponds to the JSON schema field "mapped_length".
MappedLength interface{} `msgpack:"mapped_length,omitempty"`
- // RegionID corresponds to the JSON schema field "region_id".
- RegionID interface{} `msgpack:"region_id,omitempty"`
-
- // RegionIndex corresponds to the JSON schema field "region_index".
- RegionIndex interface{} `msgpack:"region_index,omitempty"`
-
// RunID corresponds to the JSON schema field "run_id".
RunID string `msgpack:"run_id"`
@@ -1949,6 +1889,19 @@ type UpdateDagRunNote struct {
Type string `msgpack:"type,omitempty"`
}
+type Warnings []interface{}
+
+// Variable schema for responses with fields that are needed for Runtime.
+type VariableResponse struct {
+ // Key corresponds to the JSON schema field "key".
+ Key string `msgpack:"key"`
+
+ // Value corresponds to the JSON schema field "value".
+ Value interface{} `msgpack:"value"`
+}
+
+type VersionData map[string]interface{}
+
// Update the response content part of an existing Human-in-the-loop response.
type UpdateHITLDetail struct {
// ChosenOptions corresponds to the JSON schema field "chosen_options".
@@ -1983,7 +1936,16 @@ type VariableKeysResult struct {
Type string `msgpack:"type,omitempty"`
}
-type VersionData map[string]interface{}
+type VariableResult struct {
+ // Key corresponds to the JSON schema field "key".
+ Key string `msgpack:"key"`
+
+ // Type corresponds to the JSON schema field "type".
+ Type string `msgpack:"type,omitempty"`
+
+ // Value corresponds to the JSON schema field "value".
+ Value interface{} `msgpack:"value"`
+}
// One positional stub-task argument pulled from an upstream task's XCom.
type XComArgBinding struct {
@@ -2000,28 +1962,6 @@ type XComArgBinding struct {
ValueSchema *ArgValueSchema `msgpack:"value_schema,omitempty"`
}
-type Warnings []interface{}
-
-// Variable schema for responses with fields that are needed for Runtime.
-type VariableResponse struct {
- // Key corresponds to the JSON schema field "key".
- Key string `msgpack:"key"`
-
- // Value corresponds to the JSON schema field "value".
- Value interface{} `msgpack:"value"`
-}
-
-type VariableResult struct {
- // Key corresponds to the JSON schema field "key".
- Key string `msgpack:"key"`
-
- // Type corresponds to the JSON schema field "type".
- Type string `msgpack:"type,omitempty"`
-
- // Value corresponds to the JSON schema field "value".
- Value interface{} `msgpack:"value"`
-}
-
type XComCountResponse struct {
// Len corresponds to the JSON schema field "len".
Len int `msgpack:"len"`
diff --git a/go-sdk/schema/supervisor-schema.json
b/go-sdk/schema/supervisor-schema.json
index 0322fe587c7..93d3250279f 100644
--- a/go-sdk/schema/supervisor-schema.json
+++ b/go-sdk/schema/supervisor-schema.json
@@ -1736,31 +1736,6 @@
},
"DeleteXCom": {
"properties": {
- "region_id": {
- "anyOf": [
- {
- "format": "uuid",
- "type": "string"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Id"
- },
- "region_index": {
- "anyOf": [
- {
- "type": "integer"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Index"
- },
"key": {
"title": "Key",
"type": "string"
@@ -2488,31 +2463,6 @@
"GetPreviousTI": {
"description": "Request to get previous task instance.",
"properties": {
- "region_id": {
- "anyOf": [
- {
- "format": "uuid",
- "type": "string"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Id"
- },
- "region_index": {
- "anyOf": [
- {
- "type": "integer"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Index"
- },
"dag_id": {
"title": "Dag Id",
"type": "string"
@@ -2566,31 +2516,6 @@
},
"GetTICount": {
"properties": {
- "region_id": {
- "anyOf": [
- {
- "format": "uuid",
- "type": "string"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Id"
- },
- "region_index": {
- "anyOf": [
- {
- "type": "integer"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Index"
- },
"dag_id": {
"title": "Dag Id",
"type": "string"
@@ -2695,31 +2620,6 @@
},
"GetTaskBreadcrumbs": {
"properties": {
- "region_id": {
- "anyOf": [
- {
- "format": "uuid",
- "type": "string"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Id"
- },
- "region_index": {
- "anyOf": [
- {
- "type": "integer"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Index"
- },
"dag_id": {
"title": "Dag Id",
"type": "string"
@@ -2794,31 +2694,6 @@
},
"GetTaskStates": {
"properties": {
- "region_id": {
- "anyOf": [
- {
- "format": "uuid",
- "type": "string"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Id"
- },
- "region_index": {
- "anyOf": [
- {
- "type": "integer"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Index"
- },
"dag_id": {
"title": "Dag Id",
"type": "string"
@@ -2961,31 +2836,6 @@
},
"GetXCom": {
"properties": {
- "region_id": {
- "anyOf": [
- {
- "format": "uuid",
- "type": "string"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Id"
- },
- "region_index": {
- "anyOf": [
- {
- "type": "integer"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Index"
- },
"key": {
"title": "Key",
"type": "string"
@@ -3038,31 +2888,6 @@
"GetXComCount": {
"description": "Get the number of (mapped) XCom values available.",
"properties": {
- "region_id": {
- "anyOf": [
- {
- "format": "uuid",
- "type": "string"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Id"
- },
- "region_index": {
- "anyOf": [
- {
- "type": "integer"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Index"
- },
"key": {
"title": "Key",
"type": "string"
@@ -3097,31 +2922,6 @@
},
"GetXComSequenceItem": {
"properties": {
- "region_id": {
- "anyOf": [
- {
- "format": "uuid",
- "type": "string"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Id"
- },
- "region_index": {
- "anyOf": [
- {
- "type": "integer"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Index"
- },
"key": {
"title": "Key",
"type": "string"
@@ -3161,31 +2961,6 @@
},
"GetXComSequenceSlice": {
"properties": {
- "region_id": {
- "anyOf": [
- {
- "format": "uuid",
- "type": "string"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Id"
- },
- "region_index": {
- "anyOf": [
- {
- "type": "integer"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Index"
- },
"key": {
"title": "Key",
"type": "string"
@@ -4077,31 +3852,6 @@
},
"SetXCom": {
"properties": {
- "region_id": {
- "anyOf": [
- {
- "format": "uuid",
- "type": "string"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Id"
- },
- "region_index": {
- "anyOf": [
- {
- "type": "integer"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Index"
- },
"key": {
"title": "Key",
"type": "string"
diff --git
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/template_rendering.py
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/template_rendering.py
index 4c273ade904..4c968fa2d1f 100644
---
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/template_rendering.py
+++
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/template_rendering.py
@@ -73,7 +73,7 @@ def render_k8s_pod_yaml(task_instance: TaskInstance) -> dict
| None:
if AIRFLOW_V_3_4_PLUS:
workload = workloads.ExecuteTask.make(task_instance)
command_args = workload_to_command_args(workload)
- map_index = workload.ti.map_index
+ map_index = workload.key.map_index
else:
command_args = generate_pod_command_args(task_instance)
map_index = task_instance.map_index
diff --git
a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/test_template_rendering.py
b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/test_template_rendering.py
index b7ddb61542a..6c722dda35b 100644
---
a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/test_template_rendering.py
+++
b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/test_template_rendering.py
@@ -47,7 +47,7 @@ DEFAULT_DATE = timezone.datetime(2021, 9, 9)
@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Regional coordinates
require Airflow 3.4+")
@mock.patch("airflow.settings.pod_mutation_hook", autospec=True)
-def
test_loop_pod_preview_uses_the_workload_public_map_index(pod_mutation_hook,
dag_maker, session):
+def
test_loop_pod_preview_is_labelled_with_the_stored_region_index(pod_mutation_hook,
dag_maker, session):
@task_group
def body():
EmptyOperator(task_id="work")
@@ -64,7 +64,7 @@ def
test_loop_pod_preview_uses_the_workload_public_map_index(pod_mutation_hook,
pod = render_k8s_pod_yaml(ti)
- assert "map_index" not in pod["metadata"]["annotations"]
+ assert pod["metadata"]["annotations"]["map_index"] == "2"
workload = json.loads(pod["spec"]["containers"][0]["args"][-1])
assert workload["ti"]["map_index"] == -1
assert workload["ti"]["region_index"] == 2
diff --git
a/providers/edge3/src/airflow/providers/edge3/worker_api/routes/logs.py
b/providers/edge3/src/airflow/providers/edge3/worker_api/routes/logs.py
index 520000e64dd..0bb985a0638 100644
--- a/providers/edge3/src/airflow/providers/edge3/worker_api/routes/logs.py
+++ b/providers/edge3/src/airflow/providers/edge3/worker_api/routes/logs.py
@@ -20,14 +20,17 @@ from __future__ import annotations
from functools import cache
from pathlib import Path
from typing import TYPE_CHECKING, Annotated
+from uuid import UUID
from fastapi import Body, Depends, status
+from sqlalchemy import select
from airflow.api_fastapi.common.db.common import SessionDep # noqa: TC001
from airflow.api_fastapi.common.router import AirflowRouter
from airflow.api_fastapi.core_api.openapi.exceptions import
create_openapi_http_exception_doc
from airflow.models.taskinstance import TaskInstance
from airflow.providers.common.compat.sdk import TaskInstanceKey, conf
+from airflow.providers.edge3.models.edge_job import EdgeJobModel
from airflow.providers.edge3.models.edge_logs import EdgeLogsModel
from airflow.providers.edge3.worker_api.auth import
jwt_token_authorization_rest
from airflow.providers.edge3.worker_api.datamodels import PushLogsBody,
WorkerApiDocs
@@ -41,6 +44,21 @@ logs_router = AirflowRouter(tags=["Logs"], prefix="/logs")
@provide_session
def _logfile_path(task: TaskInstanceKey, *, session=NEW_SESSION) -> str:
"""Elaborate the (relative) path and filename to expect from task
execution."""
+ ti_id = session.scalar(
+ select(EdgeJobModel.task_instance_id).where(
+ EdgeJobModel.dag_id == task.dag_id,
+ EdgeJobModel.task_id == task.task_id,
+ EdgeJobModel.run_id == task.run_id,
+ EdgeJobModel.map_index == task.map_index,
+ EdgeJobModel.try_number == task.try_number,
+ EdgeJobModel.task_instance_id != "",
+ )
+ )
+ if ti_id:
+ ti = session.get(TaskInstance, UUID(ti_id),
execution_options={"include_all_attempts": True})
+ if TYPE_CHECKING:
+ assert ti
+ return FileTaskHandler(".")._render_filename(ti, task.try_number)
ti = TaskInstance.get_task_instance(
dag_id=task.dag_id,
run_id=task.run_id,
diff --git
a/providers/edge3/src/airflow/providers/edge3/worker_api/v2-edge-generated.yaml
b/providers/edge3/src/airflow/providers/edge3/worker_api/v2-edge-generated.yaml
index 9ec177fa319..827b5a802c1 100644
---
a/providers/edge3/src/airflow/providers/edge3/worker_api/v2-edge-generated.yaml
+++
b/providers/edge3/src/airflow/providers/edge3/worker_api/v2-edge-generated.yaml
@@ -1336,6 +1336,17 @@ components:
type: integer
title: Map Index
default: -1
+ region_id:
+ anyOf:
+ - type: string
+ format: uuid
+ - type: 'null'
+ title: Region Id
+ region_index:
+ anyOf:
+ - type: integer
+ - type: 'null'
+ title: Region Index
hostname:
anyOf:
- type: string
diff --git a/providers/edge3/tests/unit/edge3/worker_api/routes/test_logs.py
b/providers/edge3/tests/unit/edge3/worker_api/routes/test_logs.py
index 1be89f07649..9c44ba92167 100644
--- a/providers/edge3/tests/unit/edge3/worker_api/routes/test_logs.py
+++ b/providers/edge3/tests/unit/edge3/worker_api/routes/test_logs.py
@@ -24,12 +24,15 @@ import pytest
from sqlalchemy import delete, select
from airflow.providers.common.compat.sdk import timezone
+from airflow.providers.edge3.models.edge_job import EdgeJobModel
from airflow.providers.edge3.models.edge_logs import EdgeLogsModel
from airflow.providers.edge3.worker_api.datamodels import PushLogsBody
from airflow.providers.edge3.worker_api.routes.logs import logfile_path,
push_logs
from airflow.providers.standard.operators.empty import EmptyOperator
from airflow.utils.session import create_session
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_4_PLUS
+
if TYPE_CHECKING:
from sqlalchemy.orm import Session
@@ -57,6 +60,46 @@ class TestLogsApiRoutes:
assert str(Path(f"dag_id={DAG_ID}") / f"run_id={RUN_ID}" /
f"task_id={TASK_ID}" / "attempt=1") in p
assert "-1" not in Path(p).parts
+ @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loop passes require
Airflow 3.4+")
+ def
test_logfile_path_of_a_loop_pass_is_resolved_through_its_edge_job(self,
dag_maker, session: Session):
+ from airflow.models.dynamic_region import DynamicRegion
+ from airflow.models.taskinstance import TaskInstance
+ from airflow.sdk import task_group
+ from airflow.sdk.definitions._internal.loop import create_loop
+
+ @task_group
+ def body():
+ EmptyOperator(task_id="work")
+
+ with dag_maker("loop_dag", serialized=True):
+ create_loop(body, max_iterations=3)
+ dr = dag_maker.create_dagrun(run_id="loop_run")
+ ti = next(ti for ti in dr.task_instances if ti.task_id == "body.work")
+ region = DynamicRegion(dag_id=dr.dag_id, run_id=dr.run_id,
node_id="body")
+ session.add(region)
+ session.flush()
+ ti.region_id, ti.region_index = region.id, 2
+ session.add(
+ EdgeJobModel(
+ dag_id=ti.dag_id,
+ task_id=ti.task_id,
+ run_id=ti.run_id,
+ map_index=ti.region_index,
+ try_number=1,
+ state="running",
+ queue="default",
+ concurrency_slots=1,
+ command="",
+ task_instance_id=str(ti.id),
+ )
+ )
+ session.commit()
+ assert session.get(TaskInstance, ti.id).region_index == 2
+
+ path = logfile_path(dag_id=ti.dag_id, task_id=ti.task_id,
run_id=ti.run_id, try_number=1, map_index=2)
+
+ assert "pass=2" in path
+
def test_push_logs(self, session: Session):
log_data = PushLogsBody(
log_chunk_data="This is Lorem Ipsum log data",
log_chunk_time=timezone.utcnow()
diff --git
a/providers/elasticsearch/src/airflow/providers/elasticsearch/log/es_task_handler.py
b/providers/elasticsearch/src/airflow/providers/elasticsearch/log/es_task_handler.py
index c70e0800137..24320cece00 100644
---
a/providers/elasticsearch/src/airflow/providers/elasticsearch/log/es_task_handler.py
+++
b/providers/elasticsearch/src/airflow/providers/elasticsearch/log/es_task_handler.py
@@ -221,12 +221,13 @@ def _strip_userinfo(url: str) -> str:
def _render_log_id(log_id_template: str, ti: TaskInstance | TaskInstanceKey,
try_number: int) -> str:
+ stored_index = getattr(ti, "region_index", None)
return log_id_template.format(
dag_id=ti.dag_id,
task_id=ti.task_id,
run_id=getattr(ti, "run_id", ""),
try_number=try_number,
- map_index=getattr(ti, "map_index", ""),
+ map_index=stored_index if stored_index is not None else getattr(ti,
"map_index", ""),
)
diff --git
a/providers/elasticsearch/tests/unit/elasticsearch/log/test_es_task_handler.py
b/providers/elasticsearch/tests/unit/elasticsearch/log/test_es_task_handler.py
index 2517975349c..b0b4e39d997 100644
---
a/providers/elasticsearch/tests/unit/elasticsearch/log/test_es_task_handler.py
+++
b/providers/elasticsearch/tests/unit/elasticsearch/log/test_es_task_handler.py
@@ -556,6 +556,12 @@ class TestElasticsearchTaskHandler:
def test_render_log_id(self, ti):
assert _render_log_id(self.es_task_handler.log_id_template, ti, 1) ==
self.LOG_ID
+ def test_render_log_id_uses_the_stored_region_index_of_a_loop_pass(self):
+ public = _MockTI(map_index=-1)
+ public.region_index = 3
+ assert _render_log_id("{map_index}", public, 1) == "3"
+ assert _render_log_id("{map_index}", _MockTI(map_index=4), 1) == "4"
+
def test_clean_date(self):
clean_logical_date = _clean_date(timezone.datetime(2016, 7, 8, 9, 10,
11, 12))
assert clean_logical_date == "2016_07_08T09_10_11_000012"
diff --git
a/providers/opensearch/src/airflow/providers/opensearch/log/os_task_handler.py
b/providers/opensearch/src/airflow/providers/opensearch/log/os_task_handler.py
index 6b61b7a6b75..96c27c1d955 100644
---
a/providers/opensearch/src/airflow/providers/opensearch/log/os_task_handler.py
+++
b/providers/opensearch/src/airflow/providers/opensearch/log/os_task_handler.py
@@ -260,12 +260,13 @@ def _strip_userinfo(url: str) -> str:
def _render_log_id(
log_id_template: str, ti: TaskInstance | TaskInstanceKey | RuntimeTI,
try_number: int
) -> str:
+ stored_index = getattr(ti, "region_index", None)
return log_id_template.format(
dag_id=ti.dag_id,
task_id=ti.task_id,
run_id=getattr(ti, "run_id", ""),
try_number=try_number,
- map_index=getattr(ti, "map_index", ""),
+ map_index=stored_index if stored_index is not None else getattr(ti,
"map_index", ""),
)
diff --git
a/providers/opensearch/tests/unit/opensearch/log/test_os_task_handler.py
b/providers/opensearch/tests/unit/opensearch/log/test_os_task_handler.py
index db3f80e9c4b..71bbecd1580 100644
--- a/providers/opensearch/tests/unit/opensearch/log/test_os_task_handler.py
+++ b/providers/opensearch/tests/unit/opensearch/log/test_os_task_handler.py
@@ -606,6 +606,12 @@ class TestOpensearchTaskHandler:
self.os_task_handler.json_format = True
assert self.os_task_handler._render_log_id(ti, 1) == self.JSON_LOG_ID
+ def test_render_log_id_uses_the_stored_region_index_of_a_loop_pass(self):
+ loop_pass = _MockTI(map_index=-1)
+ loop_pass.region_index = 3
+ assert _render_log_id("{map_index}", loop_pass, 1) == "3"
+ assert _render_log_id("{map_index}", _MockTI(map_index=4), 1) == "4"
+
def test_clean_date(self):
clean_logical_date =
OpensearchTaskHandler._clean_date(timezone.datetime(2016, 7, 8, 9, 10, 11, 12))
assert clean_logical_date == "2016_07_08T09_10_11_000012"
diff --git a/task-sdk/src/airflow/sdk/api/client.py
b/task-sdk/src/airflow/sdk/api/client.py
index 4f184bf6bdf..872c11cca13 100644
--- a/task-sdk/src/airflow/sdk/api/client.py
+++ b/task-sdk/src/airflow/sdk/api/client.py
@@ -249,15 +249,6 @@ def _log_and_trace_retry(retry_state) -> None:
)
-def _region_params(region_id: uuid.UUID | None, region_index: int | None) ->
dict[str, str | int]:
- params: dict[str, str | int] = {}
- if region_id is not None:
- params["region_id"] = str(region_id)
- if region_index is not None:
- params["region_index"] = region_index
- return params
-
-
class TaskInstanceOperations:
__slots__ = ("client",)
@@ -411,9 +402,6 @@ class TaskInstanceOperations:
logical_dates: list[datetime] | None = None,
run_ids: list[str] | None = None,
states: list[str] | None = None,
- *,
- region_id: uuid.UUID | None = None,
- region_index: int | None = None,
) -> TICount:
"""Get count of task instances matching the given criteria."""
params: dict[str, Any]
@@ -432,7 +420,6 @@ class TaskInstanceOperations:
if map_index is not None and map_index >= 0:
params.update({"map_index": map_index})
- params.update(_region_params(region_id, region_index))
resp = self.client.get("task-instances/count", params=params)
return TICount(count=resp.json())
@@ -443,9 +430,6 @@ class TaskInstanceOperations:
logical_date: datetime | None = None,
map_index: int = -1,
state: TaskInstanceState | str | None = None,
- *,
- region_id: uuid.UUID | None = None,
- region_index: int | None = None,
) -> PreviousTIResult:
"""
Get the previous task instance matching the given criteria.
@@ -462,7 +446,6 @@ class TaskInstanceOperations:
if state:
params["state"] = state.value if isinstance(state,
TaskInstanceState) else state
- params.update(_region_params(region_id, region_index))
resp = self.client.get(f"task-instances/previous/{dag_id}/{task_id}",
params=params)
return PreviousTIResult(task_instance=resp.json())
@@ -474,9 +457,6 @@ class TaskInstanceOperations:
task_group_id: str | None = None,
logical_dates: list[datetime] | None = None,
run_ids: list[str] | None = None,
- *,
- region_id: uuid.UUID | None = None,
- region_index: int | None = None,
) -> TaskStatesResponse:
"""Get task states given criteria."""
params: dict[str, Any]
@@ -494,7 +474,6 @@ class TaskInstanceOperations:
if map_index is not None and map_index >= 0:
params.update({"map_index": map_index})
- params.update(_region_params(region_id, region_index))
resp = self.client.get("task-instances/states", params=params)
return TaskStatesResponse.model_validate_json(resp.read())
@@ -502,11 +481,8 @@ class TaskInstanceOperations:
self,
dag_id: str,
run_id: str,
- *,
- region_id: uuid.UUID | None = None,
- region_index: int | None = None,
) -> TaskBreadcrumbsResponse:
- params = {"dag_id": dag_id, "run_id": run_id,
**_region_params(region_id, region_index)}
+ params = {"dag_id": dag_id, "run_id": run_id}
resp = self.client.get("task-instances/breadcrumbs", params=params)
return TaskBreadcrumbsResponse.model_validate_json(resp.read())
@@ -630,14 +606,9 @@ class XComOperations:
run_id: str,
task_id: str,
key: str,
- *,
- region_id: uuid.UUID | None = None,
- region_index: int | None = None,
) -> XComCountResponse:
"""Get the number of mapped XCom values."""
- resp = self.client.head(
- f"xcoms/{dag_id}/{run_id}/{task_id}/{key}",
params=_region_params(region_id, region_index)
- )
+ resp = self.client.head(f"xcoms/{dag_id}/{run_id}/{task_id}/{key}")
# content_range: str | None
if not (content_range := resp.headers["Content-Range"]) or not
content_range.startswith(
@@ -654,12 +625,9 @@ class XComOperations:
key: str,
map_index: int | None = None,
include_prior_dates: bool = False,
- *,
- region_id: uuid.UUID | None = None,
- region_index: int | None = None,
) -> XComResponse:
"""Get a XCom value from the API server."""
- params = _region_params(region_id, region_index)
+ params: dict[str, str | int] = {}
if map_index is not None and map_index >= 0:
params.update({"map_index": map_index})
if include_prior_dates:
@@ -696,11 +664,9 @@ class XComOperations:
*,
dag_result: bool = False,
mapped_length: int | None = None,
- region_id: uuid.UUID | None = None,
- region_index: int | None = None,
) -> OKResponse:
"""Set a XCom value via the API server."""
- params = _region_params(region_id, region_index)
+ params: dict[str, str | int] = {}
if dag_result:
params["dag_result"] = dag_result
if map_index is not None and map_index >= 0:
@@ -720,12 +686,9 @@ class XComOperations:
task_id: str,
key: str,
map_index: int | None = None,
- *,
- region_id: uuid.UUID | None = None,
- region_index: int | None = None,
) -> OKResponse:
"""Delete a XCom with given key via the API server."""
- params = _region_params(region_id, region_index)
+ params: dict[str, str | int] = {}
if map_index is not None and map_index >= 0:
params["map_index"] = map_index
self.client.delete(f"xcoms/{dag_id}/{run_id}/{task_id}/{key}",
params=params)
@@ -741,15 +704,9 @@ class XComOperations:
task_id: str,
key: str,
offset: int,
- *,
- region_id: uuid.UUID | None = None,
- region_index: int | None = None,
) -> XComSequenceIndexResponse | ErrorResponse:
try:
- resp = self.client.get(
- f"xcoms/{dag_id}/{run_id}/{task_id}/{key}/item/{offset}",
- params=_region_params(region_id, region_index),
- )
+ resp =
self.client.get(f"xcoms/{dag_id}/{run_id}/{task_id}/{key}/item/{offset}")
except ServerResponseError as e:
if e.response.status_code == HTTPStatus.NOT_FOUND:
log.error(
@@ -785,11 +742,8 @@ class XComOperations:
stop: int | None,
step: int | None,
include_prior_dates: bool = False,
- *,
- region_id: uuid.UUID | None = None,
- region_index: int | None = None,
) -> XComSequenceSliceResponse:
- params = _region_params(region_id, region_index)
+ params: dict[str, str | int] = {}
if start is not None:
params["start"] = start
if stop is not None:
diff --git a/task-sdk/src/airflow/sdk/execution_time/comms.py
b/task-sdk/src/airflow/sdk/execution_time/comms.py
index e47c0463ce6..6472863fbd8 100644
--- a/task-sdk/src/airflow/sdk/execution_time/comms.py
+++ b/task-sdk/src/airflow/sdk/execution_time/comms.py
@@ -941,8 +941,6 @@ class
SkipDownstreamTasks(TISkippedDownstreamTasksStatePayload):
class GetXCom(BaseModel):
- region_id: UUID | None = None
- region_index: int | None = None
key: str
dag_id: str
run_id: str
@@ -955,8 +953,6 @@ class GetXCom(BaseModel):
class GetXComCount(BaseModel):
"""Get the number of (mapped) XCom values available."""
- region_id: UUID | None = None
- region_index: int | None = None
key: str
dag_id: str
run_id: str
@@ -965,8 +961,6 @@ class GetXComCount(BaseModel):
class GetXComSequenceItem(BaseModel):
- region_id: UUID | None = None
- region_index: int | None = None
key: str
dag_id: str
run_id: str
@@ -976,8 +970,6 @@ class GetXComSequenceItem(BaseModel):
class GetXComSequenceSlice(BaseModel):
- region_id: UUID | None = None
- region_index: int | None = None
key: str
dag_id: str
run_id: str
@@ -990,8 +982,6 @@ class GetXComSequenceSlice(BaseModel):
class SetXCom(BaseModel):
- region_id: UUID | None = None
- region_index: int | None = None
key: str
value: JsonValue
dag_id: str
@@ -1004,8 +994,6 @@ class SetXCom(BaseModel):
class DeleteXCom(BaseModel):
- region_id: UUID | None = None
- region_index: int | None = None
key: str
dag_id: str
run_id: str
@@ -1178,8 +1166,6 @@ class GetPreviousDagRun(BaseModel):
class GetPreviousTI(BaseModel):
"""Request to get previous task instance."""
- region_id: UUID | None = None
- region_index: int | None = None
dag_id: str
task_id: str
logical_date: AwareDatetime | None = None
@@ -1245,8 +1231,6 @@ class GetTaskRescheduleStartDate(BaseModel):
class GetTICount(BaseModel):
- region_id: UUID | None = None
- region_index: int | None = None
dag_id: str
map_index: int | None = None
task_ids: list[str] | None = None
@@ -1258,8 +1242,6 @@ class GetTICount(BaseModel):
class GetTaskStates(BaseModel):
- region_id: UUID | None = None
- region_index: int | None = None
dag_id: str
map_index: int | None = None
task_ids: list[str] | None = None
@@ -1270,8 +1252,6 @@ class GetTaskStates(BaseModel):
class GetTaskBreadcrumbs(BaseModel):
- region_id: UUID | None = None
- region_index: int | None = None
dag_id: str
run_id: str
type: Literal["GetTaskBreadcrumbs"] = "GetTaskBreadcrumbs"
diff --git a/task-sdk/src/airflow/sdk/execution_time/request_handlers.py
b/task-sdk/src/airflow/sdk/execution_time/request_handlers.py
index b81248cfb43..594a871f3a8 100644
--- a/task-sdk/src/airflow/sdk/execution_time/request_handlers.py
+++ b/task-sdk/src/airflow/sdk/execution_time/request_handlers.py
@@ -142,8 +142,6 @@ def handle_delete_variable(client: Client, msg:
DeleteVariable) -> tuple[BaseMod
def handle_get_ti_count(client: Client, msg: GetTICount) -> tuple[BaseModel |
None, dict[str, bool]]:
"""Fetch task instance counts."""
resp = client.task_instances.get_count(
- region_id=msg.region_id,
- region_index=msg.region_index,
dag_id=msg.dag_id,
map_index=msg.map_index,
task_ids=msg.task_ids,
@@ -158,8 +156,6 @@ def handle_get_ti_count(client: Client, msg: GetTICount) ->
tuple[BaseModel | No
def handle_get_task_states(client: Client, msg: GetTaskStates) ->
tuple[BaseModel | None, dict[str, bool]]:
"""Fetch task states and normalize them for supervisor response
handling."""
task_states_map = client.task_instances.get_task_states(
- region_id=msg.region_id,
- region_index=msg.region_index,
dag_id=msg.dag_id,
map_index=msg.map_index,
task_ids=msg.task_ids,
@@ -175,8 +171,6 @@ def handle_get_task_states(client: Client, msg:
GetTaskStates) -> tuple[BaseMode
def handle_get_previous_ti(client: Client, msg: GetPreviousTI) ->
tuple[BaseModel | None, dict[str, bool]]:
"""Fetch the previous task instance."""
resp = client.task_instances.get_previous(
- region_id=msg.region_id,
- region_index=msg.region_index,
dag_id=msg.dag_id,
task_id=msg.task_id,
logical_date=msg.logical_date,
@@ -197,8 +191,6 @@ def handle_set_xcom(client: Client, msg: SetXCom) ->
tuple[BaseModel | None, dic
msg.map_index,
dag_result=msg.dag_result,
mapped_length=msg.mapped_length,
- region_id=msg.region_id,
- region_index=msg.region_index,
)
return None, {}
@@ -211,8 +203,6 @@ def handle_delete_xcom(client: Client, msg: DeleteXCom) ->
tuple[BaseModel | Non
msg.task_id,
msg.key,
msg.map_index,
- region_id=msg.region_id,
- region_index=msg.region_index,
)
return None, {}
@@ -267,9 +257,7 @@ def handle_get_prev_successful_dag_run(
def handle_get_xcom_count(client: Client, msg: GetXComCount) ->
tuple[BaseModel | None, dict[str, bool]]:
"""Fetch XCom count metadata."""
- resp = client.xcoms.head(
- msg.dag_id, msg.run_id, msg.task_id, msg.key, region_id=msg.region_id,
region_index=msg.region_index
- )
+ resp = client.xcoms.head(msg.dag_id, msg.run_id, msg.task_id, msg.key)
return resp, {}
@@ -283,8 +271,6 @@ def handle_get_xcom_sequence_item(
msg.task_id,
msg.key,
msg.offset,
- region_id=msg.region_id,
- region_index=msg.region_index,
)
if isinstance(xcom, XComSequenceIndexResponse):
return XComSequenceIndexResult.from_response(xcom), {}
@@ -304,8 +290,6 @@ def handle_get_xcom_sequence_slice(
msg.stop,
msg.step,
msg.include_prior_dates,
- region_id=msg.region_id,
- region_index=msg.region_index,
)
if isinstance(xcoms, XComSequenceSliceResponse):
return XComSequenceSliceResult.from_response(xcoms), {}
@@ -321,8 +305,6 @@ def handle_get_xcom(client: Client, msg: GetXCom) ->
tuple[BaseModel | None, dic
msg.key,
msg.map_index,
msg.include_prior_dates,
- region_id=msg.region_id,
- region_index=msg.region_index,
)
if isinstance(xcom, XComResponse):
xcom_result = XComResult.from_xcom_response(xcom)
diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json
b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json
index 0322fe587c7..93d3250279f 100644
--- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json
+++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json
@@ -1736,31 +1736,6 @@
},
"DeleteXCom": {
"properties": {
- "region_id": {
- "anyOf": [
- {
- "format": "uuid",
- "type": "string"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Id"
- },
- "region_index": {
- "anyOf": [
- {
- "type": "integer"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Index"
- },
"key": {
"title": "Key",
"type": "string"
@@ -2488,31 +2463,6 @@
"GetPreviousTI": {
"description": "Request to get previous task instance.",
"properties": {
- "region_id": {
- "anyOf": [
- {
- "format": "uuid",
- "type": "string"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Id"
- },
- "region_index": {
- "anyOf": [
- {
- "type": "integer"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Index"
- },
"dag_id": {
"title": "Dag Id",
"type": "string"
@@ -2566,31 +2516,6 @@
},
"GetTICount": {
"properties": {
- "region_id": {
- "anyOf": [
- {
- "format": "uuid",
- "type": "string"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Id"
- },
- "region_index": {
- "anyOf": [
- {
- "type": "integer"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Index"
- },
"dag_id": {
"title": "Dag Id",
"type": "string"
@@ -2695,31 +2620,6 @@
},
"GetTaskBreadcrumbs": {
"properties": {
- "region_id": {
- "anyOf": [
- {
- "format": "uuid",
- "type": "string"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Id"
- },
- "region_index": {
- "anyOf": [
- {
- "type": "integer"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Index"
- },
"dag_id": {
"title": "Dag Id",
"type": "string"
@@ -2794,31 +2694,6 @@
},
"GetTaskStates": {
"properties": {
- "region_id": {
- "anyOf": [
- {
- "format": "uuid",
- "type": "string"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Id"
- },
- "region_index": {
- "anyOf": [
- {
- "type": "integer"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Index"
- },
"dag_id": {
"title": "Dag Id",
"type": "string"
@@ -2961,31 +2836,6 @@
},
"GetXCom": {
"properties": {
- "region_id": {
- "anyOf": [
- {
- "format": "uuid",
- "type": "string"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Id"
- },
- "region_index": {
- "anyOf": [
- {
- "type": "integer"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Index"
- },
"key": {
"title": "Key",
"type": "string"
@@ -3038,31 +2888,6 @@
"GetXComCount": {
"description": "Get the number of (mapped) XCom values available.",
"properties": {
- "region_id": {
- "anyOf": [
- {
- "format": "uuid",
- "type": "string"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Id"
- },
- "region_index": {
- "anyOf": [
- {
- "type": "integer"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Index"
- },
"key": {
"title": "Key",
"type": "string"
@@ -3097,31 +2922,6 @@
},
"GetXComSequenceItem": {
"properties": {
- "region_id": {
- "anyOf": [
- {
- "format": "uuid",
- "type": "string"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Id"
- },
- "region_index": {
- "anyOf": [
- {
- "type": "integer"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Index"
- },
"key": {
"title": "Key",
"type": "string"
@@ -3161,31 +2961,6 @@
},
"GetXComSequenceSlice": {
"properties": {
- "region_id": {
- "anyOf": [
- {
- "format": "uuid",
- "type": "string"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Id"
- },
- "region_index": {
- "anyOf": [
- {
- "type": "integer"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Index"
- },
"key": {
"title": "Key",
"type": "string"
@@ -4077,31 +3852,6 @@
},
"SetXCom": {
"properties": {
- "region_id": {
- "anyOf": [
- {
- "format": "uuid",
- "type": "string"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Id"
- },
- "region_index": {
- "anyOf": [
- {
- "type": "integer"
- },
- {
- "type": "null"
- }
- ],
- "default": null,
- "title": "Region Index"
- },
"key": {
"title": "Key",
"type": "string"
diff --git
a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py
b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py
index 5f656655216..bdbb825c2d1 100644
--- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py
+++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py
@@ -41,7 +41,6 @@ def get_bundle() -> VersionBundle:
AddArgBindingsToSupervisorTIRunContext,
AddDagDefinitionsToDagFileParsingResult,
AddRegionCoordinatesToSupervisorTaskInstance,
- AddRegionSelectors,
AddRetryReasonToTaskState,
)
@@ -53,7 +52,6 @@ def get_bundle() -> VersionBundle:
AddRetryReasonToTaskState,
AddDagDefinitionsToDagFileParsingResult,
AddRegionCoordinatesToSupervisorTaskInstance,
- AddRegionSelectors,
),
Version("2026-06-16"),
)
diff --git
a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py
b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py
index aa1e64b4e09..e093ecd07ec 100644
--- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py
+++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py
@@ -31,48 +31,10 @@ from cadwyn import VersionChange, schema
from airflow.dag_processing.processor import DagFileParsingResult # noqa:
SDK002
from airflow.sdk.api.datamodels._generated import PreviousTIResponse,
TaskInstance, TIRunContext
from airflow.sdk.execution_time.comms import (
- DeleteXCom,
- GetPreviousTI,
- GetTaskBreadcrumbs,
- GetTaskStates,
- GetTICount,
- GetXCom,
- GetXComCount,
- GetXComSequenceItem,
- GetXComSequenceSlice,
- SetXCom,
TaskState,
)
-class AddRegionSelectors(VersionChange):
- """Carry explicit coordinates in XCom and task instance requests."""
-
- description = __doc__
- instructions_to_migrate_to_previous_version = (
- schema(GetTICount).field("region_id").didnt_exist,
- schema(GetTICount).field("region_index").didnt_exist,
- schema(GetTaskStates).field("region_id").didnt_exist,
- schema(GetTaskStates).field("region_index").didnt_exist,
- schema(GetPreviousTI).field("region_id").didnt_exist,
- schema(GetPreviousTI).field("region_index").didnt_exist,
- schema(GetTaskBreadcrumbs).field("region_id").didnt_exist,
- schema(GetTaskBreadcrumbs).field("region_index").didnt_exist,
- schema(GetXCom).field("region_id").didnt_exist,
- schema(GetXCom).field("region_index").didnt_exist,
- schema(GetXComCount).field("region_id").didnt_exist,
- schema(GetXComCount).field("region_index").didnt_exist,
- schema(GetXComSequenceItem).field("region_id").didnt_exist,
- schema(GetXComSequenceItem).field("region_index").didnt_exist,
- schema(GetXComSequenceSlice).field("region_id").didnt_exist,
- schema(GetXComSequenceSlice).field("region_index").didnt_exist,
- schema(SetXCom).field("region_id").didnt_exist,
- schema(SetXCom).field("region_index").didnt_exist,
- schema(DeleteXCom).field("region_id").didnt_exist,
- schema(DeleteXCom).field("region_index").didnt_exist,
- )
-
-
class AddRegionCoordinatesToSupervisorTaskInstance(VersionChange):
"""Carry task coordinates in supervisor startup and previous-instance
messages."""
diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
index f8a963a9e3d..a5f164daa6d 100644
--- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
+++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
@@ -2216,8 +2216,6 @@ class ActivitySubprocess(WatchedSubprocess):
api_resp = self.client.task_instances.get_task_breakcrumbs(
dag_id=msg.dag_id,
run_id=msg.run_id,
- region_id=msg.region_id,
- region_index=msg.region_index,
)
resp = TaskBreadcrumbsResult.from_api_response(api_resp)
return resp, {}
diff --git a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py
b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py
index ed2ca761773..da3f120f063 100644
--- a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py
+++ b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py
@@ -52,17 +52,7 @@ from task_sdk.execution_time.schema._mock_version_bundle
import (
from airflow.sdk import TaskInstanceState
from airflow.sdk.api.datamodels._generated import PreviousTIResponse
from airflow.sdk.execution_time.comms import (
- DeleteXCom,
- GetPreviousTI,
- GetTaskBreadcrumbs,
- GetTaskStates,
- GetTICount,
- GetXCom,
- GetXComCount,
- GetXComSequenceItem,
- GetXComSequenceSlice,
PreviousTIResult,
- SetXCom,
TaskState,
)
from airflow.sdk.execution_time.schema import (
@@ -81,38 +71,6 @@ class _MockBody(BaseModel):
sentry_trace_id: str | None = None
[email protected](
- ("model", "extra"),
- [
- (GetXCom, {}),
- (GetXComCount, {}),
- (GetXComSequenceItem, {"offset": 0}),
- (GetXComSequenceSlice, {"start": None, "stop": None, "step": None}),
- (SetXCom, {"value": "value"}),
- (DeleteXCom, {}),
- (GetTICount, {}),
- (GetTaskStates, {}),
- (GetPreviousTI, {}),
- (GetTaskBreadcrumbs, {}),
- ],
-)
-def
test_region_selectors_are_versioned_and_legacy_messages_remain_unscoped(model,
extra):
- real_migrator = get_schema_version_migrator()
- msg = model(
- dag_id="dag", run_id="run", task_id="task", key="key",
region_id=UUID(int=0), region_index=-1, **extra
- )
-
- legacy = real_migrator.downgrade(msg, "2026-06-16").model_dump()
- assert "region_id" not in legacy
- assert "region_index" not in legacy
- upgraded = real_migrator.upgrade(legacy, model, "2026-06-16")
- assert upgraded["region_id"] is None
- assert upgraded["region_index"] is None
- current = real_migrator.downgrade(msg, "2026-10-30").model_dump()
- assert current["region_id"] == UUID(int=0)
- assert current["region_index"] == -1
-
-
@pytest.mark.parametrize("version", ["2026-06-16", "2026-10-30"])
def test_previous_ti_response_coordinates_follow_supervisor_version(version):
message = PreviousTIResult(
diff --git a/task-sdk/tests/task_sdk/execution_time/test_request_handlers.py
b/task-sdk/tests/task_sdk/execution_time/test_request_handlers.py
index 5a297bef111..12a31532462 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_request_handlers.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_request_handlers.py
@@ -16,14 +16,9 @@
# under the License.
from __future__ import annotations
-import json
from unittest.mock import MagicMock
-from uuid import UUID, uuid4
-import httpx
import pytest
-import structlog
-from task_sdk import make_client
from airflow.sdk.api import client as sdk_client
from airflow.sdk.api.datamodels._generated import AssetStateStoreResponse
@@ -34,42 +29,22 @@ from airflow.sdk.execution_time.comms import (
ClearAssetStateStoreByUri,
DeleteAssetStateStoreByName,
DeleteAssetStateStoreByUri,
- DeleteXCom,
ErrorResponse,
GetAssetStateStoreByName,
GetAssetStateStoreByUri,
- GetPreviousTI,
- GetTaskBreadcrumbs,
- GetTaskStates,
- GetTICount,
- GetXCom,
- GetXComCount,
- GetXComSequenceItem,
- GetXComSequenceSlice,
SetAssetStateStoreByName,
SetAssetStateStoreByUri,
- SetXCom,
)
from airflow.sdk.execution_time.request_handlers import (
handle_clear_asset_state_store_by_name,
handle_clear_asset_state_store_by_uri,
handle_delete_asset_state_store_by_name,
handle_delete_asset_state_store_by_uri,
- handle_delete_xcom,
handle_get_asset_state_store_by_name,
handle_get_asset_state_store_by_uri,
- handle_get_previous_ti,
- handle_get_task_states,
- handle_get_ti_count,
- handle_get_xcom,
- handle_get_xcom_count,
- handle_get_xcom_sequence_item,
- handle_get_xcom_sequence_slice,
handle_set_asset_state_store_by_name,
handle_set_asset_state_store_by_uri,
- handle_set_xcom,
)
-from airflow.sdk.execution_time.supervisor import ActivitySubprocess
@pytest.fixture
@@ -84,58 +59,6 @@ def client_ssl_cache():
sdk_client.Client._get_ssl_context_cached.cache_clear()
[email protected](
- ("message", "handler", "extra", "response"),
- [
- (GetXCom, handle_get_xcom, {}, {"key": "key", "value": "result"}),
- (GetXComCount, handle_get_xcom_count, {}, None),
- (GetXComSequenceItem, handle_get_xcom_sequence_item, {"offset": 0},
"result"),
- (
- GetXComSequenceSlice,
- handle_get_xcom_sequence_slice,
- {"start": 0, "stop": 2, "step": 1},
- ["result"],
- ),
- (SetXCom, handle_set_xcom, {"value": "result"}, None),
- (DeleteXCom, handle_delete_xcom, {}, None),
- (GetTICount, handle_get_ti_count, {}, 1),
- (GetTaskStates, handle_get_task_states, {}, {"task_states": {}}),
- (GetPreviousTI, handle_get_previous_ti, {}, None),
- (
- GetTaskBreadcrumbs,
- lambda client, msg:
ActivitySubprocess._handle_get_task_breadcrumbs(
- MagicMock(spec=ActivitySubprocess, client=client), msg,
structlog.get_logger(), 1
- ),
- {},
- {"breadcrumbs": []},
- ),
- ],
-)
[email protected]("coordinates", [None, (UUID(int=0), -1), (uuid4(),
2)])
-def test_region_selectors_survive_comms_and_http_transport(
- message, handler, extra, response, coordinates, client_ssl_cache
-):
- selectors = {} if coordinates is None else {"region_id": coordinates[0],
"region_index": coordinates[1]}
- msg = message(dag_id="dag", run_id="run", task_id="task", key="key",
**extra, **selectors)
- decoded = message.model_validate_json(msg.model_dump_json())
- requests = []
-
- def respond(request):
- requests.append(request)
- return httpx.Response(200, content=json.dumps(response),
headers={"Content-Range": "map_indexes 1"})
-
- handler(make_client(transport=httpx.MockTransport(respond)), decoded)
-
- assert len(requests) == 1
- params = requests[0].url.params
- if coordinates is None:
- assert "region_id" not in params
- assert "region_index" not in params
- else:
- assert params["region_id"] == str(coordinates[0])
- assert params["region_index"] == str(coordinates[1])
-
-
def test_get_asset_state_store_by_name_wraps_response_as_result(client):
client.asset_state_store.get.return_value =
AssetStateStoreResponse(value="2026-01-01")
diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
index 63aabf4b14a..d109dd4ffd7 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
@@ -2441,7 +2441,6 @@ REQUEST_TEST_CASES = [
client_mock=ClientMock(
method_path="xcoms.get",
args=("test_dag", "test_run", "test_task", "test_key", None,
False),
- kwargs={"region_id": None, "region_index": None},
response=XComResult(key="test_key", value="test_value"),
),
expected_body={"key": "test_key", "value": "test_value", "type":
"XComResult"},
@@ -2454,7 +2453,6 @@ REQUEST_TEST_CASES = [
client_mock=ClientMock(
method_path="xcoms.get",
args=("test_dag", "test_run", "test_task", "test_key", 2, False),
- kwargs={"region_id": None, "region_index": None},
response=XComResult(key="test_key", value="test_value"),
),
expected_body={"key": "test_key", "value": "test_value", "type":
"XComResult"},
@@ -2465,7 +2463,6 @@ REQUEST_TEST_CASES = [
client_mock=ClientMock(
method_path="xcoms.get",
args=("test_dag", "test_run", "test_task", "test_key", None,
False),
- kwargs={"region_id": None, "region_index": None},
response=XComResult(key="test_key", value=None, type="XComResult"),
),
expected_body={"key": "test_key", "value": None, "type": "XComResult"},
@@ -2482,7 +2479,6 @@ REQUEST_TEST_CASES = [
client_mock=ClientMock(
method_path="xcoms.get",
args=("test_dag", "test_run", "test_task", "test_key", None, True),
- kwargs={"region_id": None, "region_index": None},
response=XComResult(key="test_key", value=None, type="XComResult"),
),
expected_body={"key": "test_key", "value": None, "type": "XComResult"},
@@ -2505,7 +2501,7 @@ REQUEST_TEST_CASES = [
'{"key": "test_key", "value": {"key2": "value2"}}',
None,
),
- kwargs={"dag_result": False, "mapped_length": None, "region_id":
None, "region_index": None},
+ kwargs={"dag_result": False, "mapped_length": None},
response=OKResponse(ok=True),
),
test_id="set_xcom",
@@ -2529,7 +2525,7 @@ REQUEST_TEST_CASES = [
'{"key": "test_key", "value": {"key2": "value2"}}',
2,
),
- kwargs={"dag_result": False, "mapped_length": None, "region_id":
None, "region_index": None},
+ kwargs={"dag_result": False, "mapped_length": None},
response=OKResponse(ok=True),
),
test_id="set_xcom_with_map_index",
@@ -2554,7 +2550,7 @@ REQUEST_TEST_CASES = [
'{"key": "test_key", "value": {"key2": "value2"}}',
2,
),
- kwargs={"dag_result": False, "mapped_length": 3, "region_id":
None, "region_index": None},
+ kwargs={"dag_result": False, "mapped_length": 3},
response=OKResponse(ok=True),
),
test_id="set_xcom_with_map_index_and_mapped_length",
@@ -2578,7 +2574,7 @@ REQUEST_TEST_CASES = [
'{"key": "test_key", "value": {"key2": "value2"}}',
None,
),
- kwargs={"dag_result": True, "mapped_length": None, "region_id":
None, "region_index": None},
+ kwargs={"dag_result": True, "mapped_length": None},
response=OKResponse(ok=True),
),
test_id="set_xcom_with_dag_result",
@@ -2594,7 +2590,6 @@ REQUEST_TEST_CASES = [
client_mock=ClientMock(
method_path="xcoms.delete",
args=("test_dag", "test_run", "test_task", "test_key", 2),
- kwargs={"region_id": None, "region_index": None},
response=OKResponse(ok=True),
),
test_id="delete_xcom",
@@ -3474,8 +3469,6 @@ REQUEST_TEST_CASES = [
client_mock=ClientMock(
method_path="task_instances.get_previous",
kwargs={
- "region_id": None,
- "region_index": None,
"dag_id": "test_dag",
"task_id": "test_task",
"logical_date": timezone.parse("2024-01-15T12:00:00Z"),
@@ -3518,8 +3511,6 @@ REQUEST_TEST_CASES = [
client_mock=ClientMock(
method_path="task_instances.get_count",
kwargs={
- "region_id": None,
- "region_index": None,
"dag_id": "test_dag",
"map_index": None,
"logical_dates": None,
@@ -3556,8 +3547,6 @@ REQUEST_TEST_CASES = [
client_mock=ClientMock(
method_path="task_instances.get_task_states",
kwargs={
- "region_id": None,
- "region_index": None,
"dag_id": "test_dag",
"map_index": None,
"task_ids": None,
@@ -3581,7 +3570,6 @@ REQUEST_TEST_CASES = [
client_mock=ClientMock(
method_path="xcoms.get_sequence_item",
args=("test_dag", "test_run", "test_task", "test_key", 0),
- kwargs={"region_id": None, "region_index": None},
response=XComSequenceIndexResult(root="test_value"),
),
test_id="get_xcom_seq_item",
@@ -3598,7 +3586,6 @@ REQUEST_TEST_CASES = [
client_mock=ClientMock(
method_path="xcoms.get_sequence_item",
args=("test_dag", "test_run", "test_task", "test_key", 2),
- kwargs={"region_id": None, "region_index": None},
response=ErrorResponse(error=ErrorType.XCOM_NOT_FOUND),
),
test_id="get_xcom_seq_item_not_found",
@@ -3618,7 +3605,6 @@ REQUEST_TEST_CASES = [
client_mock=ClientMock(
method_path="xcoms.get_sequence_slice",
args=("test_dag", "test_run", "test_task", "test_key", None, None,
None, False),
- kwargs={"region_id": None, "region_index": None},
response=XComSequenceSliceResult(root=["foo", "bar"]),
),
test_id="get_xcom_seq_slice",
@@ -3681,7 +3667,6 @@ REQUEST_TEST_CASES = [
client_mock=ClientMock(
method_path="xcoms.head",
args=("test_dag", "test_run", "test_task", "test_key"),
- kwargs={"region_id": None, "region_index": None},
response=XComCountResponse(len=5),
),
test_id="get_xcom_count",
@@ -3704,7 +3689,7 @@ REQUEST_TEST_CASES = [
message=GetTaskBreadcrumbs(dag_id="test_dag", run_id="test_run"),
client_mock=ClientMock(
method_path="task_instances.get_task_breakcrumbs",
- kwargs={"dag_id": "test_dag", "run_id": "test_run", "region_id":
None, "region_index": None},
+ kwargs={"dag_id": "test_dag", "run_id": "test_run"},
response=TaskBreadcrumbsResult(
breadcrumbs=[
{
diff --git a/ts-sdk/src/generated/supervisor.ts
b/ts-sdk/src/generated/supervisor.ts
index ba766af0f48..3878f891ed6 100644
--- a/ts-sdk/src/generated/supervisor.ts
+++ b/ts-sdk/src/generated/supervisor.ts
@@ -339,8 +339,6 @@ export type Key3 = string;
export type Type23 = "DeleteTaskStateStore";
export type Key4 = string;
export type Type24 = "DeleteVariable";
-export type RegionId1 = string | null;
-export type RegionIndex1 = number | null;
export type Key5 = string;
export type DagId6 = string;
export type RunId5 = string;
@@ -440,15 +438,11 @@ export type DagId11 = string;
export type LogicalDate3 = string;
export type State3 = string | null;
export type Type41 = "GetPreviousDagRun";
-export type RegionId2 = string | null;
-export type RegionIndex2 = number | null;
export type DagId12 = string;
export type TaskId3 = string;
export type LogicalDate4 = string | null;
export type MapIndex2 = number;
export type Type42 = "GetPreviousTI";
-export type RegionId3 = string | null;
-export type RegionIndex3 = number | null;
export type DagId13 = string;
export type MapIndex3 = number | null;
export type TaskIds = string[] | null;
@@ -457,8 +451,6 @@ export type LogicalDates1 = string[] | null;
export type RunIds1 = string[] | null;
export type States1 = string[] | null;
export type Type43 = "GetTICount";
-export type RegionId4 = string | null;
-export type RegionIndex4 = number | null;
export type DagId14 = string;
export type RunId8 = string;
export type Type44 = "GetTaskBreadcrumbs";
@@ -468,8 +460,6 @@ export type Type45 = "GetTaskRescheduleStartDate";
export type TiId6 = string;
export type Key8 = string;
export type Type46 = "GetTaskStateStore";
-export type RegionId5 = string | null;
-export type RegionIndex5 = number | null;
export type DagId15 = string;
export type MapIndex4 = number | null;
export type TaskIds1 = string[] | null;
@@ -483,8 +473,6 @@ export type Prefix = string | null;
export type Limit2 = number;
export type Offset = number;
export type Type49 = "GetVariableKeys";
-export type RegionId6 = string | null;
-export type RegionIndex6 = number | null;
export type Key10 = string;
export type DagId16 = string;
export type RunId9 = string;
@@ -492,23 +480,17 @@ export type TaskId4 = string;
export type MapIndex5 = number | null;
export type IncludePriorDates = boolean;
export type Type50 = "GetXCom";
-export type RegionId7 = string | null;
-export type RegionIndex7 = number | null;
export type Key11 = string;
export type DagId17 = string;
export type RunId10 = string;
export type TaskId5 = string;
export type Type51 = "GetXComCount";
-export type RegionId8 = string | null;
-export type RegionIndex8 = number | null;
export type Key12 = string;
export type DagId18 = string;
export type RunId11 = string;
export type TaskId6 = string;
export type Offset1 = number;
export type Type52 = "GetXComSequenceItem";
-export type RegionId9 = string | null;
-export type RegionIndex9 = number | null;
export type Key13 = string;
export type DagId19 = string;
export type RunId12 = string;
@@ -553,8 +535,8 @@ export type EndDate4 = string | null;
export type State4 = string | null;
export type TryNumber2 = number;
export type MapIndex6 = number | null;
-export type RegionId10 = string | null;
-export type RegionIndex10 = number | null;
+export type RegionId1 = string | null;
+export type RegionIndex1 = number | null;
export type Duration = number | null;
export type Type60 = "PreviousTIResult";
export type Key14 = string;
@@ -587,8 +569,6 @@ export type TiId8 = string;
export type Key17 = string;
export type ExpiresAt = string | null;
export type Type70 = "SetTaskStateStore";
-export type RegionId11 = string | null;
-export type RegionIndex11 = number | null;
export type Key18 = string;
export type DagId21 = string;
export type RunId14 = string;
@@ -1279,8 +1259,6 @@ export interface DeleteVariable {
* via the `definition` "DeleteXCom".
*/
export interface DeleteXCom {
- region_id?: RegionId1;
- region_index?: RegionIndex1;
key: Key5;
dag_id: DagId6;
run_id: RunId5;
@@ -1450,8 +1428,6 @@ export interface GetPreviousDagRun {
* via the `definition` "GetPreviousTI".
*/
export interface GetPreviousTI {
- region_id?: RegionId2;
- region_index?: RegionIndex2;
dag_id: DagId12;
task_id: TaskId3;
logical_date?: LogicalDate4;
@@ -1464,8 +1440,6 @@ export interface GetPreviousTI {
* via the `definition` "GetTICount".
*/
export interface GetTICount {
- region_id?: RegionId3;
- region_index?: RegionIndex3;
dag_id: DagId13;
map_index?: MapIndex3;
task_ids?: TaskIds;
@@ -1480,8 +1454,6 @@ export interface GetTICount {
* via the `definition` "GetTaskBreadcrumbs".
*/
export interface GetTaskBreadcrumbs {
- region_id?: RegionId4;
- region_index?: RegionIndex4;
dag_id: DagId14;
run_id: RunId8;
type?: Type44;
@@ -1509,8 +1481,6 @@ export interface GetTaskStateStore {
* via the `definition` "GetTaskStates".
*/
export interface GetTaskStates {
- region_id?: RegionId5;
- region_index?: RegionIndex5;
dag_id: DagId15;
map_index?: MapIndex4;
task_ids?: TaskIds1;
@@ -1542,8 +1512,6 @@ export interface GetVariableKeys {
* via the `definition` "GetXCom".
*/
export interface GetXCom {
- region_id?: RegionId6;
- region_index?: RegionIndex6;
key: Key10;
dag_id: DagId16;
run_id: RunId9;
@@ -1559,8 +1527,6 @@ export interface GetXCom {
* via the `definition` "GetXComCount".
*/
export interface GetXComCount {
- region_id?: RegionId7;
- region_index?: RegionIndex7;
key: Key11;
dag_id: DagId17;
run_id: RunId10;
@@ -1572,8 +1538,6 @@ export interface GetXComCount {
* via the `definition` "GetXComSequenceItem".
*/
export interface GetXComSequenceItem {
- region_id?: RegionId8;
- region_index?: RegionIndex8;
key: Key12;
dag_id: DagId18;
run_id: RunId11;
@@ -1586,8 +1550,6 @@ export interface GetXComSequenceItem {
* via the `definition` "GetXComSequenceSlice".
*/
export interface GetXComSequenceSlice {
- region_id?: RegionId9;
- region_index?: RegionIndex9;
key: Key13;
dag_id: DagId19;
run_id: RunId12;
@@ -1681,8 +1643,8 @@ export interface PreviousTIResponse {
state?: State4;
try_number: TryNumber2;
map_index?: MapIndex6;
- region_id?: RegionId10;
- region_index?: RegionIndex10;
+ region_id?: RegionId1;
+ region_index?: RegionIndex1;
duration?: Duration;
}
/**
@@ -1805,8 +1767,6 @@ export interface SetTaskStateStore {
* via the `definition` "SetXCom".
*/
export interface SetXCom {
- region_id?: RegionId11;
- region_index?: RegionIndex11;
key: Key18;
value: JsonValue;
dag_id: DagId21;
