This is an automated email from the ASF dual-hosted git repository.
vatsrahul1001 pushed a commit to branch v3-3-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-3-test by this push:
new 2dbf4cf35dc [v3-3-test] Reuse the explicit dag_run join on the
previous-TI lookup (#72811) (#72944)
2dbf4cf35dc is described below
commit 2dbf4cf35dcdb6c9e2693ffd890cc925afda5acc
Author: github-actions[bot]
<41898282+github-actions[bot]@users.noreply.github.com>
AuthorDate: Fri Sep 11 17:24:38 2026 +0530
[v3-3-test] Reuse the explicit dag_run join on the previous-TI lookup
(#72811) (#72944)
* Reuse the explicit dag_run join on the previous-TI lookup
The previous-TI query in ``get_previous_task_instance`` explicitly joins
``dag_run`` for the ``ORDER BY dag_run.logical_date`` and asked for a second
``joinedload(TI.dag_run)`` for hydration. ``TaskInstance.dag_run`` is
already
``lazy="joined"`` on the mapper, so the eager-load actually attached a
parallel
``JOIN dag_run AS dag_run_1`` next to the explicit join. Both joins pull the
same rows over the same FK pair.
Switch to ``contains_eager(TI.dag_run)`` so the eager-load reuses the
explicit
join instead of adding a second one, and add a regression that fails when
the
compiled SQL references ``dag_run`` more than once in the FROM clause.
* Drop unused load_only import
* Assert the eager-load claim by watching for stray dag_run SELECTs
* Swap the stray-select capture for a request-level query count
(cherry picked from commit d076a57dcb24303c69d8505ea4816417003ffb34)
Co-authored-by: Pierre Jeambrun <[email protected]>
---
.../api_fastapi/execution_api/routes/task_instances.py | 4 ++--
.../execution_api/versions/head/test_task_instances.py | 18 ++++++++++++++----
2 files changed, 16 insertions(+), 6 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 f8b8d7e7de1..a7eef46571b 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
@@ -36,7 +36,7 @@ from pydantic import JsonValue
from sqlalchemy import and_, func, or_, tuple_, update
from sqlalchemy.engine import CursorResult
from sqlalchemy.exc import DataError, NoResultFound, SQLAlchemyError
-from sqlalchemy.orm import joinedload
+from sqlalchemy.orm import contains_eager, joinedload
from sqlalchemy.sql import select
from structlog.contextvars import bind_contextvars
@@ -1194,7 +1194,7 @@ def get_previous_task_instance(
query = (
select(TI)
.join(DR, (TI.dag_id == DR.dag_id) & (TI.run_id == DR.run_id))
- .options(joinedload(TI.dag_run))
+ .options(contains_eager(TI.dag_run).load_only(DR.logical_date))
.where(TI.dag_id == dag_id, TI.task_id == task_id, TI.map_index ==
map_index)
.order_by(DR.logical_date.desc())
)
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 7d4d5cc7433..9796399f8ca 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
@@ -58,7 +58,7 @@ from airflow.sdk import Asset, TaskGroup, TriggerRule, task,
task_group
from airflow.state.metastore import MetastoreBackend
from airflow.utils.state import DagRunState, State, TaskInstanceState,
TerminalTIState
-from tests_common.test_utils.asserts import capture_orm_selects
+from tests_common.test_utils.asserts import assert_queries_count,
capture_orm_selects
from tests_common.test_utils.config import conf_vars
from tests_common.test_utils.db import (
clear_db_assets,
@@ -3747,7 +3747,8 @@ class TestGetPreviousTI:
assert data["state"] == State.SUCCESS
def test_get_previous_ti_query_is_bounded(self, client, session,
create_task_instance):
- """The single-row previous-TI lookup must ask the DB for one row."""
+ """The single-row previous-TI lookup must ask the DB for one row, join
``dag_run`` once,
+ and surface the eager-loaded ``logical_date`` in the response."""
for i in range(5):
create_task_instance(
task_id="test_task",
@@ -3757,17 +3758,26 @@ class TestGetPreviousTI:
)
session.commit()
- with capture_orm_selects("task_instance") as statements:
+ with (
+ capture_orm_selects("task_instance") as statements,
+ assert_queries_count(1),
+ ):
response = client.get(
"/execution/task-instances/previous/dag/test_task",
params={"logical_date": "2025-01-05T00:00:00Z"},
)
assert response.status_code == 200
- assert response.json()["run_id"] == "run4"
+ data = response.json()
+ assert data["run_id"] == "run4"
+ assert data["logical_date"] == "2025-01-04T00:00:00Z"
assert statements, "expected the endpoint to query the task_instance
table"
for sql in statements:
assert re.search(r"\bLIMIT 1\b", sql), f"previous-TI lookup is not
bounded to one row: {sql}"
+ dag_run_join_count = len(re.findall(r"JOIN dag_run(\s|$)", sql))
+ assert dag_run_join_count == 1, (
+ f"previous-TI query joins dag_run {dag_run_join_count} times,
expected once: {sql}"
+ )
class TestGetTaskStates:
