OscarLigthart commented on code in PR #62812:
URL: https://github.com/apache/airflow/pull/62812#discussion_r3116700513
##########
airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py:
##########
@@ -101,14 +112,72 @@ def _patch_ti_validate_request(
return dag, list(tis), body.model_dump(include=fields_to_update,
by_alias=True)
+def _get_task_group_task_instances(
+ dag_id: str,
+ dag_run_id: str,
+ task_group_id: str,
+ dag: SerializedDAG,
+ session: Session,
+) -> list[TI]:
+ """Get all task instances in a task group for a specific DAG run."""
+ task_group = dag.task_group_dict.get(task_group_id)
+ if not task_group:
+ raise HTTPException(
+ status.HTTP_404_NOT_FOUND, f"Task group '{task_group_id}' not
found in DAG '{dag_id}'"
+ )
+
+ task_ids = [task.task_id for task in task_group.iter_tasks()]
+
+ query = (
+ select(TI)
+ .where(
+ TI.dag_id == dag_id,
+ TI.run_id == dag_run_id,
+ TI.task_id.in_(task_ids),
+ )
+ .join(TI.dag_run)
+ .options(joinedload(TI.rendered_task_instance_fields))
+ .order_by(TI.task_id, TI.map_index)
+ )
+
+ group_tis = list(session.scalars(query).all())
+
+ return group_tis
+
+
+def _patch_ti_group_validate_request(
+ dag_id: str,
+ dag_run_id: str,
+ task_group_id: str,
+ dag_bag: DagBagDep,
+ body: PatchTaskInstanceBody,
+ session: SessionDep,
+ update_mask: list[str] | None = Query(None),
+) -> tuple[SerializedDAG, list[TI], dict]:
+ """Validate and prepare data for task group patch request."""
+ dag = get_latest_version_of_dag(dag_bag, dag_id, session)
+ tis = _get_task_group_task_instances(dag_id, dag_run_id, task_group_id,
dag, session)
+
+ fields_to_update = body.model_fields_set
+ if update_mask:
+ fields_to_update = fields_to_update.intersection(update_mask)
+ else:
+ try:
+ PatchTaskInstanceBody.model_validate(body)
+ except ValidationError as e:
+ raise RequestValidationError(errors=e.errors())
+
+ return dag, tis, body.model_dump(include=fields_to_update, by_alias=True)
Review Comment:
I took it out of both to prevent duplicate logic, but I'm not entirely sure
if that's what you meant. Let me know 👍
--
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]
