dheerajturaga commented on code in PR #56837:
URL: https://github.com/apache/airflow/pull/56837#discussion_r2443540199
##########
airflow-core/src/airflow/api_fastapi/core_api/services/ui/grid.py:
##########
@@ -130,3 +139,148 @@ def _find_aggregates(
**_get_aggs_for_node(details),
}
return
+
+
+def get_batch_ti_summaries(
+ dag_id: str,
+ run_ids: list[str],
+ session: Session,
+) -> dict:
+ """
+ Fetch task instance summaries for multiple runs in a single query.
+
+ This is much more efficient than the N+1 query pattern of calling
+ get_grid_ti_summaries multiple times.
+
+ Returns a dict with structure:
+ {
+ "dag_id": str,
+ "summaries": [GridTISummaries, ...]
+ }
+ """
+ if not run_ids:
+ return {"dag_id": dag_id, "summaries": []}
+
+ # Single query to fetch ALL task instances for ALL runs
+ tis_query = (
+ select(
+ TaskInstance.run_id,
+ TaskInstance.task_id,
+ TaskInstance.state,
+ TaskInstance.dag_version_id,
+ TaskInstance.start_date,
+ TaskInstance.end_date,
+ )
+ .where(
+ TaskInstance.dag_id == dag_id,
+ TaskInstance.run_id.in_(run_ids),
+ )
+ .order_by(TaskInstance.run_id, TaskInstance.task_id)
+ )
+
+ task_instances = list(session.execute(tis_query))
+
+ # Group by run_id and collect unique dag_version_ids
+ tis_by_run = collections.defaultdict(list)
+ dag_version_ids = set()
+ for ti in task_instances:
+ tis_by_run[ti.run_id].append(ti)
+ if ti.dag_version_id:
+ dag_version_ids.add(ti.dag_version_id)
+
+ # Fetch all needed serialized DAGs in one query
+ serdags_by_version = {}
+ if dag_version_ids:
+ serdags = session.scalars(
+ select(SerializedDagModel)
+ .join(DagVersion, SerializedDagModel.dag_version_id ==
DagVersion.id)
+ .where(DagVersion.id.in_(dag_version_ids))
+ )
+ for serdag in serdags:
+ if serdag.dag_version_id:
+ serdags_by_version[serdag.dag_version_id] = serdag
Review Comment:
@jason810496 , Can you elaborate more on what you would prefer to be done
here?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
