[GitHub] [airflow] bbovenzi commented on a change in pull request #22314: WIP summarize mapped

[GitBox <gi...@apache.org>]

bbovenzi commented on a change in pull request #22314:
URL: https://github.com/apache/airflow/pull/22314#discussion_r828232501



##########
File path: airflow/api_connexion/endpoints/task_instance_endpoint.py
##########
@@ -296,10 +298,49 @@ def get_task_instances_batch(session: Session = NEW_SESSION) -> APIResponse:
     ti_query = base_query.options(joinedload(TI.rendered_task_instance_fields))
     task_instances = ti_query.all()
 
-    return task_instance_collection_schema.dump(
+    results = task_instance_collection_schema.dump(
         TaskInstanceCollection(task_instances=task_instances, total_entries=total_entries)
     )
 
+    if "summarize_mapped" in body and body["summarize_mapped"]:
+        dag_run_ids = [ti["dag_run_id"] for ti in results["task_instances"]]
+        mapped_ti_query = session.query(TI).join(TI.dag_run)
+        mapped_ti_query = _apply_array_filter(mapped_ti_query, key=TI.run_id, values=dag_run_ids)
+        mapped_ti_query = mapped_ti_query.filter(TI.map_index != -1)
+        # FIXME without SLA block, this error when rendering:
+        #       TypeError: 'TaskInstance' object is not subscriptable
+        mapped_ti_query = mapped_ti_query.join(
+            SlaMiss,
+            and_(
+                SlaMiss.dag_id == TI.dag_id,
+                SlaMiss.task_id == TI.task_id,
+                SlaMiss.execution_date == DR.execution_date,
+            ),
+            isouter=True,
+        ).add_entity(SlaMiss)
+        mapped_ti_query = mapped_ti_query.options(joinedload(TI.rendered_task_instance_fields))
+        mapped_task_instances = mapped_ti_query.all()
+        mapped_summaries = task_instance_summary_collection_schema.dump(
+            TaskInstanceCollection(task_instances=mapped_task_instances, total_entries=1)
+        )
+
+        by_dag_run_id = {}
+        for mapped_ti in mapped_summaries["task_instances"]:
+            dag_run_id = mapped_ti["dag_run_id"]
+            try:
+                by_dag_run_id[dag_run_id].append(mapped_ti)
+            except:
+                by_dag_run_id[dag_run_id] = [
+                    mapped_ti,
+                ]
+
+        for ti in results["task_instances"]:
+            dag_run_id = ti["dag_run_id"]
+            if dag_run_id in by_dag_run_id:
+                ti["mapped_tasks"] = by_dag_run_id[dag_run_id]

Review comment:
       Let's have the list of `mapped_tasks` be a separate endpoint with which we can paginate with `limit` and `offset` params (and possibly, search).




-- 
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: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org