From 686424e248371314996fca4f22f592e201053d57 Mon Sep 17 00:00:00 2001 From: ColtenOuO Date: Sat, 29 Aug 2026 18:32:40 +0000 Subject: [PATCH] Remove redundant full-Dag TaskInstance query from task-states execution API get_task_instance_states always ran a full, unfiltered TaskInstance query for the Dag even when scoped to a single task_group_id, then discarded those rows in favor of a second, separate query for the group's members. Workers call this endpoint on the task execution hot path, so the wasted round trip added avoidable RTT and DB load on every task-group-scoped call, proportional to Dag size and Dag-run count. --- .../execution_api/routes/task_instances.py | 62 ++++++++++--------- .../versions/head/test_task_instances.py | 34 ++++++++++ 2 files changed, 67 insertions(+), 29 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 3dd4bc11a0bfd..88b44a44496ae 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 @@ -1271,9 +1271,13 @@ def get_task_instance_states( """Get the states for Task Instances with the given criteria.""" run_id_task_state_map: dict[str, dict[str, Any]] = defaultdict(dict) - query = select(TI).where(TI.dag_id == dag_id) + query = select(TI.run_id, TI.task_id, TI.map_index, TI.state).where(TI.dag_id == dag_id) - if task_ids: + if task_group_id: + group_task_ids = _get_group_task_ids(dag_id, task_group_id, session, dag_bag) + effective_task_ids = set(task_ids) | group_task_ids if task_ids else group_task_ids + query = query.where(TI.task_id.in_(effective_task_ids)) + elif task_ids: query = query.where(TI.task_id.in_(task_ids)) if logical_dates: @@ -1285,23 +1289,9 @@ def get_task_instance_states( if map_index is not None: query = query.where(TI.map_index == map_index) - results = session.scalars(query).all() - - if task_group_id: - group_tasks = _get_group_tasks( - dag_id, task_group_id, session, dag_bag, logical_dates, run_ids, map_index - ) - - results = results + group_tasks if task_ids else group_tasks - - [ - run_id_task_state_map[task.run_id].update( - {task.task_id: task.state} - if task.map_index < 0 - else {f"{task.task_id}_{task.map_index}": task.state} - ) - for task in results - ] + for ti in session.execute(query): + key = ti.task_id if ti.map_index < 0 else f"{ti.task_id}_{ti.map_index}" + run_id_task_state_map[ti.run_id][key] = ti.state return TaskStatesResponse(task_states=run_id_task_state_map) @@ -1333,16 +1323,7 @@ def _is_eligible_to_retry(state: str, try_number: int, max_tries: int) -> bool: return max_tries != 0 and try_number <= max_tries -def _get_group_tasks( - dag_id: str, - task_group_id: str, - session: SessionDep, - dag_bag: DagBagDep, - logical_dates=None, - run_ids=None, - map_index: int | None = None, -): - # Get all tasks in the task group +def _get_task_group(dag_id: str, task_group_id: str, session: SessionDep, dag_bag: DagBagDep): dag = get_latest_version_of_dag(dag_bag, dag_id, session, include_reason=True) task_group = dag.task_group_dict.get(task_group_id) if not task_group: @@ -1353,6 +1334,29 @@ def _get_group_tasks( "message": f"Task group {task_group_id} not found in DAG {dag_id}", }, ) + return task_group + + +def _get_group_task_ids( + dag_id: str, + task_group_id: str, + session: SessionDep, + dag_bag: DagBagDep, +) -> set[str]: + task_group = _get_task_group(dag_id, task_group_id, session, dag_bag) + return {task.task_id for task in task_group.iter_tasks()} + + +def _get_group_tasks( + dag_id: str, + task_group_id: str, + session: SessionDep, + dag_bag: DagBagDep, + logical_dates=None, + run_ids=None, + map_index: int | None = None, +): + task_group = _get_task_group(dag_id, task_group_id, session, dag_bag) # First get all task instances to get the task_id, map_index pairs group_tasks = session.scalars( 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 8ecf7b1b5dd17..ba4f1ec4c2848 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 @@ -4112,6 +4112,40 @@ def test_get_task_states_with_task_group_id_and_task_id(self, client, session, d }, } + def test_get_task_states_with_task_group_id_query_count(self, client, session, dag_maker): + """Regression test: a task_group_id lookup must not also run a full, unfiltered TI query for the Dag.""" + with dag_maker("test_get_task_states_with_task_group_id_query_count", serialized=True): + with TaskGroup("group1"): + EmptyOperator(task_id="task1") + EmptyOperator(task_id="task2") + EmptyOperator(task_id="task3") + + dr = dag_maker.create_dagrun() + + tis = dr.get_task_instances() + for ti in tis: + ti.state = State.SUCCESS + session.merge(ti) + session.commit() + + with assert_queries_count(2): + response = client.get( + "/execution/task-instances/states", + params={ + "dag_id": "test_get_task_states_with_task_group_id_query_count", + "task_group_id": "group1", + }, + ) + assert response.status_code == 200 + assert response.json() == { + "task_states": { + "test": { + "group1.task1": "success", + "group1.task2": "success", + }, + }, + } + def test_get_task_group_states_with_multiple_task(self, client, session, dag_maker): with dag_maker("test_get_task_group_states_with_multiple_task_tasks", serialized=True): with TaskGroup("group1"):