Skip to content

Commit 171524f

Browse files
hkc-8010eladkal
authored andcommitted
Reduce reschedule startup queries in execution API
The run endpoint was doing two TaskReschedule round trips to compute values that can be fetched together. Collapsing them into one aggregate query keeps the rescheduled-sensor context fix intact while trimming overhead on task startup, and the test now proves we return the earliest reschedule timestamp instead of depending on insertion order.
1 parent d539cee commit 171524f

2 files changed

Lines changed: 12 additions & 17 deletions

File tree

airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py

Lines changed: 6 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -290,17 +290,12 @@ def ti_run(
290290
xcom_query = xcom_query.where(XComModel.map_index == map_index)
291291

292292
xcom_keys = list(session.scalars(xcom_query))
293-
task_reschedule_count = (
294-
session.scalar(
295-
select(func.count(TaskReschedule.id)).where(TaskReschedule.ti_id == task_instance_id)
296-
)
297-
or 0
298-
)
299-
first_task_reschedule_start_date = None
300-
if task_reschedule_count > 0:
301-
first_task_reschedule_start_date = session.scalar(
302-
select(func.min(TaskReschedule.start_date)).where(TaskReschedule.ti_id == task_instance_id)
303-
)
293+
task_reschedule_count, first_task_reschedule_start_date = session.execute(
294+
select(
295+
func.count(TaskReschedule.id),
296+
func.min(TaskReschedule.start_date),
297+
).where(TaskReschedule.ti_id == task_instance_id)
298+
).one()
304299

305300
dr.team_name = get_team_name_for_ti(task_instance_id, session)
306301

airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -346,18 +346,18 @@ def test_ti_run_state_includes_first_task_reschedule_start_date(
346346
)
347347
session.add_all(
348348
[
349-
TaskReschedule(
350-
ti_id=ti.id,
351-
start_date=first_reschedule_start_date,
352-
end_date=timezone.datetime(2024, 9, 30, 10, 1),
353-
reschedule_date=timezone.datetime(2024, 9, 30, 10, 2),
354-
),
355349
TaskReschedule(
356350
ti_id=ti.id,
357351
start_date=second_reschedule_start_date,
358352
end_date=timezone.datetime(2024, 9, 30, 11, 1),
359353
reschedule_date=timezone.datetime(2024, 9, 30, 11, 2),
360354
),
355+
TaskReschedule(
356+
ti_id=ti.id,
357+
start_date=first_reschedule_start_date,
358+
end_date=timezone.datetime(2024, 9, 30, 10, 1),
359+
reschedule_date=timezone.datetime(2024, 9, 30, 10, 2),
360+
),
361361
]
362362
)
363363
session.commit()

0 commit comments

Comments
 (0)