pierrejeambrun commented on code in PR #62812:
URL: https://github.com/apache/airflow/pull/62812#discussion_r3111265581
##########
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))
Review Comment:
`joinedload(TI.rendered_task_instance_fields)` is applied here AND
re-queried with `populate_existing=True` in the route after
`set_task_group_state` returns. One of them is wasted — probably this one since
the route always re-queries. Can we drop it here?
##########
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)
Review Comment:
Is this join needed? The `dag_id` + `run_id` filters already constrain the
result. I see `_patch_ti_validate_request` does the same thing — pre-existing
but worth cleaning up in both places while we're here.
##########
airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py:
##########
@@ -56,6 +56,17 @@
log = structlog.get_logger(__name__)
+def _collect_unique_tis(
+ affected_tis_dict: dict[tuple[str, str, str, int], TI],
+ tis: list[TI] | None,
+) -> None:
+ """Collect unique task instances into a dictionary keyed by (dag_id,
run_id, task_id, map_index)."""
+ if tis:
+ for ti in tis:
+ key = (ti.dag_id, ti.run_id, ti.task_id, ti.map_index)
+ affected_tis_dict[key] = ti
Review Comment:
Do we really need this? `set_task_group_state` already returns unique TIs
(the underlying `set_state` dedupes), and the `note` branch only touches TIs
already in the group. Feels defensive for no real case.
##########
airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py:
##########
@@ -6477,3 +6478,383 @@ def test_should_respond_403(self,
unauthorized_test_client):
def test_should_respond_422(self, test_client):
response = test_client.patch(self.ENDPOINT_URL, json={})
assert response.status_code == 422
+
+
+class TestPatchTaskGroup(TestTaskInstanceEndpoint):
Review Comment:
Most tests in this class mock `set_task_group_state`, which means we're only
validating the route/service glue. The real v2.x-equivalent logic we just added
in `SerializedDAG.set_task_group_state` isn't exercised. Can we add unmocked
tests for:
- `include_upstream=True` / `include_downstream=True` actually affecting
up/downstream TIs
- `past=True` / `future=True` covering the right runs (and the
`exclude_run_ids` logic)
- The `clear()` of failed downstream actually resuming them
- The 409 path (all TIs already in target state)
The existing `test_patch_task_group_updates_ti_states_in_db` and the
query-count tests are exactly the right pattern — just more of those.
##########
airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py:
##########
@@ -145,6 +214,53 @@ def _patch_task_instance_state(
except Exception:
log.exception("error calling listener")
+ return updated_tis
+
+
+def _patch_task_group_state(
+ group_id: str,
+ dag_run_id: str,
+ dag: SerializedDAG,
+ body: PatchTaskInstanceBody,
+ data: dict,
+ *,
+ session: Session,
+) -> list[TI]:
+ """Update the state of all task instances in a task group."""
+ updated_tis = dag.set_task_group_state(
+ group_id=group_id,
+ run_id=dag_run_id,
+ state=data["new_state"],
+ upstream=body.include_upstream,
+ downstream=body.include_downstream,
+ future=body.include_future,
+ past=body.include_past,
+ commit=True,
Review Comment:
Nit: `commit=True` is hardcoded here but the helper could easily take it as
a param. Not critical, but would make this reusable for a dry-run path later
without duplicating the function.
##########
airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml:
##########
@@ -7348,6 +7348,166 @@ paths:
application/json:
schema:
$ref: '#/components/schemas/HTTPValidationError'
+ /api/v2/dags/{dag_id}/dagRuns/{dag_run_id}/taskGroupInstances/{group_id}:
+ patch:
+ tags:
+ - Task Instance
+ summary: Patch Task Group Instances
+ description: Update the state of all task instances in a task group.
+ operationId: patch_task_group_instances
Review Comment:
Can you double-check the generated spec for both endpoints? Especially that
the 409 response is registered for the main patch endpoint (since we raise it
when all TIs are already in the target state).
##########
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:
This function is basically a copy of `_patch_ti_validate_request` with
`task_id`/`map_index` swapped for `task_group_id` and a different fetch helper.
Can we factor out the common part (fields_to_update + validation at the bottom)?
##########
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),
Review Comment:
`Query(None)` is misleading here — this helper isn't a FastAPI dep, it's
called internally. Just use `update_mask: list[str] | None = None`. Same issue
in `_patch_ti_validate_request` above, worth fixing both.
--
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]
