From 478a173ce5f6e4a7f8e600dfa0df57c6a6e296ab Mon Sep 17 00:00:00 2001 From: Shivam Rastogi <6463385+shivaam@users.noreply.github.com> Date: Sun, 22 Mar 2026 22:45:20 -0700 Subject: [PATCH 1/2] feat: Add dagrun status progress to backfills --- .../core_api/datamodels/backfills.py | 2 + .../core_api/openapi/_private_ui.yaml | 10 +++ .../openapi/v2-rest-api-generated.yaml | 10 +++ .../core_api/routes/public/backfills.py | 63 ++++++++++++++++--- .../core_api/routes/ui/backfills.py | 6 +- .../ui/openapi-gen/requests/schemas.gen.ts | 13 ++++ .../ui/openapi-gen/requests/types.gen.ts | 4 ++ .../ui/public/i18n/locales/en/common.json | 2 + .../ui/src/components/BackfillProgressBar.tsx | 63 +++++++++++++++++++ .../src/components/Banner/BackfillBanner.tsx | 10 ++- .../ui/src/pages/Dag/Backfills/Backfills.tsx | 28 +++++++++ .../core_api/routes/public/test_backfills.py | 60 +++++++++++++++++- .../core_api/routes/ui/test_backfills.py | 8 ++- .../airflowctl/api/datamodels/generated.py | 2 + 14 files changed, 267 insertions(+), 14 deletions(-) create mode 100644 airflow-core/src/airflow/ui/src/components/BackfillProgressBar.tsx 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..d309228bd7d8b 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] = {} 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..1ef1f4b393ff0 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 @@ -16,13 +16,14 @@ # under the License. from __future__ import annotations +from collections import defaultdict from typing import Annotated from fastapi import Depends, HTTPException, status from fastapi.exceptions import RequestValidationError from pydantic import NonNegativeInt -from sqlalchemy import select, update -from sqlalchemy.orm import joinedload +from sqlalchemy import func, select, update +from sqlalchemy.orm import Session, joinedload from airflow._shared.timezones import timezone from airflow.api_fastapi.common.db.common import ( @@ -61,6 +62,40 @@ backfills_router = AirflowRouter(tags=["Backfill"], prefix="/backfills") +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, + ) + .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][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 = counts.get(backfill.id, {}) + return backfills + + @backfills_router.get( path="", dependencies=[ @@ -84,8 +119,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 +141,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 +172,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 +200,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 +253,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..4a95979901c11 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 @@ -36,6 +36,7 @@ from airflow.api_fastapi.core_api.openapi.exceptions import ( create_openapi_http_exception_doc, ) +from airflow.api_fastapi.core_api.routes.public.backfills import _enrich_backfill_responses from airflow.api_fastapi.core_api.security import ReadableBackfillsFilterDep, requires_access_backfill from airflow.models.backfill import Backfill @@ -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/src/airflow/ui/public/i18n/locales/en/common.json b/airflow-core/src/airflow/ui/public/i18n/locales/en/common.json index e44f9cc3856e1..3cd0a0929ed49 100644 --- a/airflow-core/src/airflow/ui/public/i18n/locales/en/common.json +++ b/airflow-core/src/airflow/ui/public/i18n/locales/en/common.json @@ -28,6 +28,7 @@ }, "collapseAllExtra": "Collapse all extra JSON", "collapseDetailsPanel": "Collapse Details Panel", + "completed": "Completed", "consumingAsset": "Consuming Asset", "createdAssetEvent_one": "Created Asset Event", "createdAssetEvent_other": "Created Asset Events", @@ -241,6 +242,7 @@ "from": "From", "maxActiveRuns": "Max Active Runs", "noTagsFound": "No tags found", + "progress": "Progress", "tagMode": { "all": "All", "any": "Any" diff --git a/airflow-core/src/airflow/ui/src/components/BackfillProgressBar.tsx b/airflow-core/src/airflow/ui/src/components/BackfillProgressBar.tsx new file mode 100644 index 0000000000000..2032887d02544 --- /dev/null +++ b/airflow-core/src/airflow/ui/src/components/BackfillProgressBar.tsx @@ -0,0 +1,63 @@ +/*! + * 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. + */ +import { Box, HStack, Text } from "@chakra-ui/react"; + +type Props = { + readonly stateCounts: Record; + readonly total: number; + readonly trackColor?: string; +}; + +export const BackfillProgressBar = ({ stateCounts, total, trackColor = "bg.emphasized" }: Props) => { + const successCount = stateCounts.success ?? 0; + const failedCount = stateCounts.failed ?? 0; + const successPct = (successCount / total) * 100; + const failedPct = (failedCount / total) * 100; + const remainingPct = 100 - successPct - failedPct; + + return ( + + + {successPct > 0 ? ( + + ) : undefined} + {failedPct > 0 ? ( + + ) : undefined} + {remainingPct > 0 ? ( + + ) : undefined} + + + {successCount + failedCount}/{total} + + + ); +}; diff --git a/airflow-core/src/airflow/ui/src/components/Banner/BackfillBanner.tsx b/airflow-core/src/airflow/ui/src/components/Banner/BackfillBanner.tsx index ddcf4d873f721..1b5e9058246e8 100644 --- a/airflow-core/src/airflow/ui/src/components/Banner/BackfillBanner.tsx +++ b/airflow-core/src/airflow/ui/src/components/Banner/BackfillBanner.tsx @@ -32,8 +32,8 @@ import { import type { BackfillResponse } from "openapi/requests/types.gen"; import { useAutoRefresh } from "src/utils"; +import { BackfillProgressBar } from "../BackfillProgressBar"; import Time from "../Time"; -import { ProgressBar } from "../ui"; type Props = { readonly dagId: string; @@ -113,7 +113,13 @@ const BackfillBanner = ({ dagId }: Props) => { - + {(backfill.num_runs ?? 0) > 0 ? ( + + ) : undefined}