diff --git a/airflow-core/src/airflow/api_fastapi/common/db/backfills.py b/airflow-core/src/airflow/api_fastapi/common/db/backfills.py new file mode 100644 index 0000000000000..2105c01971778 --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/common/db/backfills.py @@ -0,0 +1,64 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from collections import defaultdict +from typing import TYPE_CHECKING + +from sqlalchemy import func, select + +from airflow.api_fastapi.core_api.datamodels.backfills import BackfillResponse +from airflow.models import DagRun +from airflow.models.backfill import BackfillDagRun + +if TYPE_CHECKING: + from sqlalchemy.orm import Session + + +def enrich_backfill_responses( + backfills: list[BackfillResponse], + *, + session: Session, +) -> list[BackfillResponse]: + """Populate num_runs and dag_run_state_counts on each backfill response.""" + ids = [b.id for b in backfills] + if not ids: + return backfills + # Single query: get state counts per backfill, derive num_runs by summing counts. + rows = session.execute( + select( + BackfillDagRun.backfill_id, + DagRun.state, + func.count().label("count"), + ) + .join(DagRun, BackfillDagRun.dag_run_id == DagRun.id) + .where( + BackfillDagRun.backfill_id.in_(ids), + DagRun.backfill_id == BackfillDagRun.backfill_id, + DagRun.state.is_not(None), + ) + .group_by(BackfillDagRun.backfill_id, DagRun.state) + ).all() + counts: dict[int, dict[str, int]] = defaultdict(dict) + num_runs: dict[int, int] = defaultdict(int) + for backfill_id, state, count in rows: + counts[backfill_id][str(state)] = count + num_runs[backfill_id] += count + for backfill in backfills: + backfill.num_runs = num_runs.get(backfill.id, 0) + backfill.dag_run_state_counts = dict(counts.get(backfill.id, {})) + return backfills diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/backfills.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/backfills.py index 87d2677028eab..af3e4285aadb9 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/backfills.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/backfills.py @@ -54,6 +54,8 @@ class BackfillResponse(BaseModel): completed_at: datetime | None updated_at: datetime dag_display_name: str = Field(validation_alias=AliasPath("dag_model", "dag_display_name")) + num_runs: int = 0 + dag_run_state_counts: dict[str, int] = Field(default_factory=dict) class BackfillCollectionResponse(BaseModel): diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml index c1553542f037c..16ebabe42fd16 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml +++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml @@ -1643,6 +1643,16 @@ components: dag_display_name: type: string title: Dag Display Name + num_runs: + type: integer + title: Num Runs + default: 0 + dag_run_state_counts: + additionalProperties: + type: integer + type: object + title: Dag Run State Counts + default: {} type: object required: - id diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml index bd8e6aa62ba45..b4e872e44e327 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml +++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml @@ -9321,6 +9321,16 @@ components: dag_display_name: type: string title: Dag Display Name + num_runs: + type: integer + title: Num Runs + default: 0 + dag_run_state_counts: + additionalProperties: + type: integer + type: object + title: Dag Run State Counts + default: {} type: object required: - id diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/backfills.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/backfills.py index c1e90ab0fa62b..5bed2cd7f893f 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/backfills.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/backfills.py @@ -25,6 +25,7 @@ from sqlalchemy.orm import joinedload from airflow._shared.timezones import timezone +from airflow.api_fastapi.common.db.backfills import enrich_backfill_responses from airflow.api_fastapi.common.db.common import ( SessionDep, paginated_select, @@ -84,8 +85,10 @@ def list_backfills( limit=limit, session=session, ) + backfills = [BackfillResponse.model_validate(b) for b in session.scalars(select_stmt)] + enrich_backfill_responses(backfills, session=session) return BackfillCollectionResponse( - backfills=session.scalars(select_stmt), + backfills=backfills, total_entries=total_entries, ) @@ -104,9 +107,11 @@ def get_backfill( backfill = session.scalars( select(Backfill).where(Backfill.id == backfill_id).options(joinedload(Backfill.dag_model)) ).one_or_none() - if backfill: - return backfill - raise HTTPException(status.HTTP_404_NOT_FOUND, "Backfill not found") + if not backfill: + raise HTTPException(status.HTTP_404_NOT_FOUND, "Backfill not found") + response = BackfillResponse.model_validate(backfill) + enrich_backfill_responses([response], session=session) + return response @backfills_router.put( @@ -133,7 +138,9 @@ def pause_backfill(backfill_id: NonNegativeInt, session: SessionDep) -> Backfill if b.is_paused is False: b.is_paused = True session.commit() - return b + response = BackfillResponse.model_validate(b) + enrich_backfill_responses([response], session=session) + return response @backfills_router.put( @@ -159,7 +166,9 @@ def unpause_backfill(backfill_id: NonNegativeInt, session: SessionDep) -> Backfi raise HTTPException(status.HTTP_409_CONFLICT, "Backfill is already completed.") if b.is_paused: b.is_paused = False - return b + response = BackfillResponse.model_validate(b) + enrich_backfill_responses([response], session=session) + return response @backfills_router.put( @@ -210,7 +219,9 @@ def cancel_backfill(backfill_id: NonNegativeInt, session: SessionDep) -> Backfil # this is in separate transaction just to avoid potential conflicts session.refresh(b) b.completed_at = timezone.utcnow() - return b + response = BackfillResponse.model_validate(b) + enrich_backfill_responses([response], session=session) + return response @backfills_router.post( diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/backfills.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/backfills.py index 02583a8355b9d..0e1016680c9dc 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/backfills.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/backfills.py @@ -22,6 +22,7 @@ from sqlalchemy import select from sqlalchemy.orm import joinedload +from airflow.api_fastapi.common.db.backfills import enrich_backfill_responses from airflow.api_fastapi.common.db.common import SessionDep, paginated_select from airflow.api_fastapi.common.parameters import ( FilterOptionEnum, @@ -73,9 +74,12 @@ def list_backfills_ui( session=session, ) backfills = [ - BackfillResponse(**row._mapping) if not isinstance(row, Backfill) else row + BackfillResponse(**row._mapping) + if not isinstance(row, Backfill) + else BackfillResponse.model_validate(row) for row in session.scalars(select_stmt) ] + enrich_backfill_responses(backfills, session=session) return BackfillCollectionResponse( backfills=backfills, total_entries=total_entries, diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts index dcb743c34f884..7bc1a59248047 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts @@ -541,6 +541,19 @@ export const $BackfillResponse = { dag_display_name: { type: 'string', title: 'Dag Display Name' + }, + num_runs: { + type: 'integer', + title: 'Num Runs', + default: 0 + }, + dag_run_state_counts: { + additionalProperties: { + type: 'integer' + }, + type: 'object', + title: 'Dag Run State Counts', + default: {} } }, type: 'object', diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts index 9d1380698f97d..d6eb444010c73 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts @@ -147,6 +147,10 @@ export type BackfillResponse = { completed_at: string | null; updated_at: string; dag_display_name: string; + num_runs?: number; + dag_run_state_counts?: { + [key: string]: (number); + }; }; /** diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_backfills.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_backfills.py index 10bfd0f781732..325f7955864d2 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_backfills.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_backfills.py @@ -126,7 +126,7 @@ def test_list_backfill(self, test_client, session): session.add(b) session.commit() - with assert_queries_count(3): + with assert_queries_count(4): response = test_client.get(f"/backfills?dag_id={dag.dag_id}") assert response.status_code == 200 @@ -145,6 +145,8 @@ def test_list_backfill(self, test_client, session): "max_active_runs": 10, "to_date": to_iso(to_date), "updated_at": mock.ANY, + "num_runs": 0, + "dag_run_state_counts": {}, } ], "total_entries": 1, @@ -174,6 +176,54 @@ def test_get_backfill(self, session, test_client): "max_active_runs": 10, "to_date": to_iso(to_date), "updated_at": mock.ANY, + "num_runs": 0, + "dag_run_state_counts": {}, + } + + def test_get_backfill_with_dag_run_state_counts(self, session, test_client): + (dag,) = self._create_dag_models() + from_date = timezone.utcnow() + to_date = timezone.utcnow() + backfill = Backfill(dag_id=dag.dag_id, from_date=from_date, to_date=to_date) + session.add(backfill) + session.flush() + + # Create dag runs associated with this backfill via BackfillDagRun + dag_runs_data = [ + (DagRunState.SUCCESS, pendulum.datetime(2024, 1, 1)), + (DagRunState.SUCCESS, pendulum.datetime(2024, 1, 2)), + (DagRunState.FAILED, pendulum.datetime(2024, 1, 3)), + (DagRunState.RUNNING, pendulum.datetime(2024, 1, 4)), + (DagRunState.QUEUED, pendulum.datetime(2024, 1, 5)), + ] + for i, (state, logical_date) in enumerate(dag_runs_data): + dr = DagRun( + dag_id=dag.dag_id, + run_id=f"backfill__{logical_date.isoformat()}", + logical_date=logical_date, + state=state, + run_type="backfill", + backfill_id=backfill.id, + ) + session.add(dr) + session.flush() + bdr = BackfillDagRun( + backfill_id=backfill.id, + dag_run_id=dr.id, + logical_date=logical_date, + sort_ordinal=i + 1, + ) + session.add(bdr) + session.commit() + + response = test_client.get(f"/backfills/{backfill.id}") + assert response.status_code == 200 + assert response.json()["num_runs"] == 5 + assert response.json()["dag_run_state_counts"] == { + "success": 2, + "failed": 1, + "running": 1, + "queued": 1, } def test_get_backfill_with_null_conf(self, session, test_client): @@ -252,6 +302,8 @@ def test_create_backfill(self, repro_act, repro_exp, session, dag_maker, test_cl "max_active_runs": 5, "to_date": to_date_iso, "updated_at": mock.ANY, + "num_runs": 0, + "dag_run_state_counts": {}, } check_last_log(session, dag_id="TEST_DAG_1", event="create_backfill", logical_date=None) @@ -830,6 +882,8 @@ def test_cancel_backfill(self, session, test_client): "max_active_runs": 10, "to_date": to_iso(to_date), "updated_at": mock.ANY, + "num_runs": 0, + "dag_run_state_counts": {}, } assert pendulum.parse(response.json()["completed_at"]) # now it is marked as completed @@ -910,6 +964,8 @@ def test_pause_backfill(self, session, test_client): "max_active_runs": 10, "to_date": to_iso(to_date), "updated_at": mock.ANY, + "num_runs": 0, + "dag_run_state_counts": {}, } check_last_log(session, dag_id=None, event="pause_backfill", logical_date=None) @@ -969,6 +1025,8 @@ def test_unpause_backfill(self, session, test_client): "max_active_runs": 10, "to_date": to_iso(to_date), "updated_at": mock.ANY, + "num_runs": 0, + "dag_run_state_counts": {}, } check_last_log(session, dag_id=None, event="unpause_backfill", logical_date=None) diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_backfills.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_backfills.py index 21ae10d23b3a9..186b8681bc293 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_backfills.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_backfills.py @@ -120,6 +120,8 @@ def test_should_response_200( "max_active_runs": 10, "to_date": from_datetime_to_zulu(to_date), "updated_at": mock.ANY, + "num_runs": 0, + "dag_run_state_counts": {}, }, "backfill2": { "completed_at": None, @@ -134,6 +136,8 @@ def test_should_response_200( "max_active_runs": 10, "to_date": from_datetime_to_zulu(to_date), "updated_at": mock.ANY, + "num_runs": 0, + "dag_run_state_counts": {}, }, "backfill3": { "completed_at": None, @@ -148,12 +152,14 @@ def test_should_response_200( "max_active_runs": 10, "to_date": from_datetime_to_zulu(to_date), "updated_at": mock.ANY, + "num_runs": 0, + "dag_run_state_counts": {}, }, } expected_response = [] for backfill in response_params: expected_response.append(backfill_responses[backfill]) - with assert_queries_count(3 if test_params.get("dag_id") is None else 4): + with assert_queries_count(4 if test_params.get("dag_id") is None else 5): response = test_client.get("/backfills", params=test_params) assert response.status_code == 200 assert response.json() == { diff --git a/airflow-ctl/src/airflowctl/api/datamodels/generated.py b/airflow-ctl/src/airflowctl/api/datamodels/generated.py index a0b32b7fbdd39..958b194cccaff 100644 --- a/airflow-ctl/src/airflowctl/api/datamodels/generated.py +++ b/airflow-ctl/src/airflowctl/api/datamodels/generated.py @@ -1155,6 +1155,8 @@ class BackfillResponse(BaseModel): completed_at: Annotated[datetime | None, Field(title="Completed At")] = None updated_at: Annotated[datetime, Field(title="Updated At")] dag_display_name: Annotated[str, Field(title="Dag Display Name")] + num_runs: Annotated[int | None, Field(title="Num Runs")] = 0 + dag_run_state_counts: Annotated[dict[str, int] | None, Field(title="Dag Run State Counts")] = {} class BulkCreateActionConnectionBody(BaseModel):