Skip to content

Commit cb0a135

Browse files
yuseok89yuseok89
committed
Show task state counts for the latest Dag run on the Dags list page
1 parent 493312b commit cb0a135

19 files changed

Lines changed: 788 additions & 10 deletions

File tree

‎airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/dags.py‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,25 @@ class DAGRunStateCountsResponse(BaseModel):
4848
state_counts: dict[DagRunState, int]
4949

5050

51+
class DAGLatestRunTaskInstanceStateCountsResponse(BaseModel):
52+
"""
53+
Task-instance state counts for a Dag's latest run.
54+
55+
``state_counts`` only carries states present in the run; task instances without a
56+
state yet are keyed as ``no_status``.
57+
"""
58+
59+
dag_id: str
60+
run_id: str
61+
state_counts: dict[str, int]
62+
63+
64+
class DAGsLatestRunTaskInstanceStateCountsCollectionResponse(BaseModel):
65+
"""Collection of per-Dag latest-run task-instance state counts for the Dag list page."""
66+
67+
dags: list[DAGLatestRunTaskInstanceStateCountsResponse]
68+
69+
5170
class DAGsRunStateCountsCollectionResponse(BaseModel):
5271
"""Collection of per-Dag DagRun-state counts for the Dag list page."""
5372

‎airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -640,6 +640,44 @@ paths:
640640
application/json:
641641
schema:
642642
$ref: '#/components/schemas/HTTPValidationError'
643+
/ui/dags/latest_run_task_instance_state_counts:
644+
get:
645+
tags:
646+
- DAG
647+
summary: Get Latest Run Task Instance State Counts
648+
description: 'Return task-instance state counts for each Dag''s latest run,
649+
for the Dag list page.
650+
651+
652+
Dags without any run are omitted from the response.'
653+
operationId: get_latest_run_task_instance_state_counts_ui
654+
security:
655+
- OAuth2PasswordBearer: []
656+
- HTTPBearer: []
657+
parameters:
658+
- name: dag_ids
659+
in: query
660+
required: true
661+
schema:
662+
type: array
663+
items:
664+
type: string
665+
minItems: 1
666+
maxItems: 100
667+
title: Dag Ids
668+
responses:
669+
'200':
670+
description: Successful Response
671+
content:
672+
application/json:
673+
schema:
674+
$ref: '#/components/schemas/DAGsLatestRunTaskInstanceStateCountsCollectionResponse'
675+
'422':
676+
description: Validation Error
677+
content:
678+
application/json:
679+
schema:
680+
$ref: '#/components/schemas/HTTPValidationError'
643681
/ui/dependencies:
644682
get:
645683
tags:
@@ -2366,6 +2404,32 @@ components:
23662404
that
23672405
23682406
the API server/Web UI can use this data to render connection form UI.'
2407+
DAGLatestRunTaskInstanceStateCountsResponse:
2408+
properties:
2409+
dag_id:
2410+
type: string
2411+
title: Dag Id
2412+
run_id:
2413+
type: string
2414+
title: Run Id
2415+
state_counts:
2416+
additionalProperties:
2417+
type: integer
2418+
type: object
2419+
title: State Counts
2420+
type: object
2421+
required:
2422+
- dag_id
2423+
- run_id
2424+
- state_counts
2425+
title: DAGLatestRunTaskInstanceStateCountsResponse
2426+
description: 'Task-instance state counts for a Dag''s latest run.
2427+
2428+
2429+
``state_counts`` only carries states present in the run; task instances without
2430+
a
2431+
2432+
state yet are keyed as ``no_status``.'
23692433
DAGRunLightResponse:
23702434
properties:
23712435
id:
@@ -2677,6 +2741,19 @@ components:
26772741
- file_token
26782742
title: DAGWithLatestDagRunsResponse
26792743
description: DAG with latest dag runs response serializer.
2744+
DAGsLatestRunTaskInstanceStateCountsCollectionResponse:
2745+
properties:
2746+
dags:
2747+
items:
2748+
$ref: '#/components/schemas/DAGLatestRunTaskInstanceStateCountsResponse'
2749+
type: array
2750+
title: Dags
2751+
type: object
2752+
required:
2753+
- dags
2754+
title: DAGsLatestRunTaskInstanceStateCountsCollectionResponse
2755+
description: Collection of per-Dag latest-run task-instance state counts for
2756+
the Dag list page.
26802757
DAGsRunStateCountsCollectionResponse:
26812758
properties:
26822759
dags:

‎airflow-core/src/airflow/api_fastapi/core_api/routes/ui/dags.py‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,9 @@
5858
from airflow.api_fastapi.core_api.datamodels.dags import DAG_ALIAS_MAPPING, DAGResponse
5959
from airflow.api_fastapi.core_api.datamodels.ui.dag_runs import DAGRunLightResponse
6060
from airflow.api_fastapi.core_api.datamodels.ui.dags import (
61+
DAGLatestRunTaskInstanceStateCountsResponse,
6162
DAGRunStateCountsResponse,
63+
DAGsLatestRunTaskInstanceStateCountsCollectionResponse,
6264
DAGsRunStateCountsCollectionResponse,
6365
DAGWithLatestDagRunsCollectionResponse,
6466
DAGWithLatestDagRunsResponse,
@@ -343,3 +345,75 @@ def get_dag_run_state_counts(
343345
],
344346
state_count_limit=STATE_COUNT_CAP,
345347
)
348+
349+
350+
@dags_router.get(
351+
"/latest_run_task_instance_state_counts",
352+
dependencies=[
353+
Depends(requires_access_dag(method="GET")),
354+
Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE)),
355+
],
356+
operation_id="get_latest_run_task_instance_state_counts_ui",
357+
)
358+
def get_latest_run_task_instance_state_counts(
359+
session: SessionDep,
360+
readable_dags_filter: ReadableDagsFilterDep,
361+
dag_ids: Annotated[list[str], Query(min_length=1, max_length=conf.getint("api", "maximum_page_limit"))],
362+
) -> DAGsLatestRunTaskInstanceStateCountsCollectionResponse:
363+
"""
364+
Return task-instance state counts for each Dag's latest run, for the Dag list page.
365+
366+
Dags without any run are omitted from the response.
367+
"""
368+
permitted_dag_ids = readable_dags_filter.value or set()
369+
requested_dag_ids = sorted(set(dag_ids) & permitted_dag_ids)
370+
371+
dags: list[DAGLatestRunTaskInstanceStateCountsResponse] = []
372+
if not requested_dag_ids:
373+
return DAGsLatestRunTaskInstanceStateCountsCollectionResponse(dags=dags)
374+
375+
latest_run_branches = [
376+
select(DagRun.dag_id, DagRun.run_id)
377+
.where(DagRun.dag_id == dag_id)
378+
.order_by(DagRun.run_after.desc())
379+
.limit(1)
380+
.subquery()
381+
for dag_id in requested_dag_ids
382+
]
383+
latest_runs_union = union_all(*(select(branch) for branch in latest_run_branches)).subquery()
384+
latest_run_id_by_dag: dict[str, str] = {
385+
row.dag_id: row.run_id for row in session.execute(select(latest_runs_union))
386+
}
387+
388+
if latest_run_id_by_dag:
389+
# Each branch filters on (dag_id, run_id) equality, which the ti_dag_run index
390+
# covers. A run's task instances are bounded by the Dag's task structure, so the
391+
# per-state counts are exact (no cap needed here, unlike the cross-run counts in
392+
# get_dag_run_state_counts).
393+
ti_branches = [
394+
select(literal(dag_id).label("dag_id"), TaskInstance.state.label("state"))
395+
.where(TaskInstance.dag_id == dag_id, TaskInstance.run_id == run_id)
396+
.subquery()
397+
for dag_id, run_id in latest_run_id_by_dag.items()
398+
]
399+
tis_union = union_all(*(select(branch) for branch in ti_branches)).subquery()
400+
counts_by_dag: dict[str, dict[str, int]] = {dag_id: {} for dag_id in latest_run_id_by_dag}
401+
for row in session.execute(
402+
select(tis_union.c.dag_id, tis_union.c.state, func.count().label("cnt")).group_by(
403+
tis_union.c.dag_id, tis_union.c.state
404+
)
405+
):
406+
state_key = row.state if row.state is not None else "no_status"
407+
counts_by_dag[row.dag_id][state_key] = row.cnt
408+
409+
dags = [
410+
DAGLatestRunTaskInstanceStateCountsResponse(
411+
dag_id=dag_id,
412+
run_id=latest_run_id_by_dag[dag_id],
413+
state_counts=counts_by_dag[dag_id],
414+
)
415+
for dag_id in requested_dag_ids
416+
if dag_id in latest_run_id_by_dag
417+
]
418+
419+
return DAGsLatestRunTaskInstanceStateCountsCollectionResponse(dags=dags)

‎airflow-core/src/airflow/ui/openapi-gen/queries/common.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -364,6 +364,12 @@ export const useDagServiceGetDagRunStateCountsUiKey = "DagServiceGetDagRunStateC
364364
export const UseDagServiceGetDagRunStateCountsUiKeyFn = ({ dagIds }: {
365365
dagIds: string[];
366366
}, queryKey?: Array<unknown>) => [useDagServiceGetDagRunStateCountsUiKey, ...(queryKey ?? [{ dagIds }])];
367+
export type DagServiceGetLatestRunTaskInstanceStateCountsUiDefaultResponse = Awaited<ReturnType<typeof DagService.getLatestRunTaskInstanceStateCountsUi>>;
368+
export type DagServiceGetLatestRunTaskInstanceStateCountsUiQueryResult<TData = DagServiceGetLatestRunTaskInstanceStateCountsUiDefaultResponse, TError = unknown> = UseQueryResult<TData, TError>;
369+
export const useDagServiceGetLatestRunTaskInstanceStateCountsUiKey = "DagServiceGetLatestRunTaskInstanceStateCountsUi";
370+
export const UseDagServiceGetLatestRunTaskInstanceStateCountsUiKeyFn = ({ dagIds }: {
371+
dagIds: string[];
372+
}, queryKey?: Array<unknown>) => [useDagServiceGetLatestRunTaskInstanceStateCountsUiKey, ...(queryKey ?? [{ dagIds }])];
367373
export type EventLogServiceGetEventLogDefaultResponse = Awaited<ReturnType<typeof EventLogService.getEventLog>>;
368374
export type EventLogServiceGetEventLogQueryResult<TData = EventLogServiceGetEventLogDefaultResponse, TError = unknown> = UseQueryResult<TData, TError>;
369375
export const useEventLogServiceGetEventLogKey = "EventLogServiceGetEventLog";

‎airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -724,6 +724,19 @@ export const ensureUseDagServiceGetDagRunStateCountsUiData = (queryClient: Query
724724
dagIds: string[];
725725
}) => queryClient.ensureQueryData({ queryKey: Common.UseDagServiceGetDagRunStateCountsUiKeyFn({ dagIds }), queryFn: () => DagService.getDagRunStateCountsUi({ dagIds }) });
726726
/**
727+
* Get Latest Run Task Instance State Counts
728+
* Return task-instance state counts for each Dag's latest run, for the Dag list page.
729+
*
730+
* Dags without any run are omitted from the response.
731+
* @param data The data for the request.
732+
* @param data.dagIds
733+
* @returns DAGsLatestRunTaskInstanceStateCountsCollectionResponse Successful Response
734+
* @throws ApiError
735+
*/
736+
export const ensureUseDagServiceGetLatestRunTaskInstanceStateCountsUiData = (queryClient: QueryClient, { dagIds }: {
737+
dagIds: string[];
738+
}) => queryClient.ensureQueryData({ queryKey: Common.UseDagServiceGetLatestRunTaskInstanceStateCountsUiKeyFn({ dagIds }), queryFn: () => DagService.getLatestRunTaskInstanceStateCountsUi({ dagIds }) });
739+
/**
727740
* Get Event Log
728741
* @param data The data for the request.
729742
* @param data.eventLogId

‎airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -724,6 +724,19 @@ export const prefetchUseDagServiceGetDagRunStateCountsUi = (queryClient: QueryCl
724724
dagIds: string[];
725725
}) => queryClient.prefetchQuery({ queryKey: Common.UseDagServiceGetDagRunStateCountsUiKeyFn({ dagIds }), queryFn: () => DagService.getDagRunStateCountsUi({ dagIds }) });
726726
/**
727+
* Get Latest Run Task Instance State Counts
728+
* Return task-instance state counts for each Dag's latest run, for the Dag list page.
729+
*
730+
* Dags without any run are omitted from the response.
731+
* @param data The data for the request.
732+
* @param data.dagIds
733+
* @returns DAGsLatestRunTaskInstanceStateCountsCollectionResponse Successful Response
734+
* @throws ApiError
735+
*/
736+
export const prefetchUseDagServiceGetLatestRunTaskInstanceStateCountsUi = (queryClient: QueryClient, { dagIds }: {
737+
dagIds: string[];
738+
}) => queryClient.prefetchQuery({ queryKey: Common.UseDagServiceGetLatestRunTaskInstanceStateCountsUiKeyFn({ dagIds }), queryFn: () => DagService.getLatestRunTaskInstanceStateCountsUi({ dagIds }) });
739+
/**
727740
* Get Event Log
728741
* @param data The data for the request.
729742
* @param data.eventLogId

‎airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -724,6 +724,19 @@ export const useDagServiceGetDagRunStateCountsUi = <TData = Common.DagServiceGet
724724
dagIds: string[];
725725
}, queryKey?: TQueryKey, options?: Omit<UseQueryOptions<TData, TError>, "queryKey" | "queryFn">) => useQuery<TData, TError>({ queryKey: Common.UseDagServiceGetDagRunStateCountsUiKeyFn({ dagIds }, queryKey), queryFn: () => DagService.getDagRunStateCountsUi({ dagIds }) as TData, ...options });
726726
/**
727+
* Get Latest Run Task Instance State Counts
728+
* Return task-instance state counts for each Dag's latest run, for the Dag list page.
729+
*
730+
* Dags without any run are omitted from the response.
731+
* @param data The data for the request.
732+
* @param data.dagIds
733+
* @returns DAGsLatestRunTaskInstanceStateCountsCollectionResponse Successful Response
734+
* @throws ApiError
735+
*/
736+
export const useDagServiceGetLatestRunTaskInstanceStateCountsUi = <TData = Common.DagServiceGetLatestRunTaskInstanceStateCountsUiDefaultResponse, TError = unknown, TQueryKey extends Array<unknown> = unknown[]>({ dagIds }: {
737+
dagIds: string[];
738+
}, queryKey?: TQueryKey, options?: Omit<UseQueryOptions<TData, TError>, "queryKey" | "queryFn">) => useQuery<TData, TError>({ queryKey: Common.UseDagServiceGetLatestRunTaskInstanceStateCountsUiKeyFn({ dagIds }, queryKey), queryFn: () => DagService.getLatestRunTaskInstanceStateCountsUi({ dagIds }) as TData, ...options });
739+
/**
727740
* Get Event Log
728741
* @param data The data for the request.
729742
* @param data.eventLogId

‎airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -724,6 +724,19 @@ export const useDagServiceGetDagRunStateCountsUiSuspense = <TData = Common.DagSe
724724
dagIds: string[];
725725
}, queryKey?: TQueryKey, options?: Omit<UseQueryOptions<TData, TError>, "queryKey" | "queryFn">) => useSuspenseQuery<TData, TError>({ queryKey: Common.UseDagServiceGetDagRunStateCountsUiKeyFn({ dagIds }, queryKey), queryFn: () => DagService.getDagRunStateCountsUi({ dagIds }) as TData, ...options });
726726
/**
727+
* Get Latest Run Task Instance State Counts
728+
* Return task-instance state counts for each Dag's latest run, for the Dag list page.
729+
*
730+
* Dags without any run are omitted from the response.
731+
* @param data The data for the request.
732+
* @param data.dagIds
733+
* @returns DAGsLatestRunTaskInstanceStateCountsCollectionResponse Successful Response
734+
* @throws ApiError
735+
*/
736+
export const useDagServiceGetLatestRunTaskInstanceStateCountsUiSuspense = <TData = Common.DagServiceGetLatestRunTaskInstanceStateCountsUiDefaultResponse, TError = unknown, TQueryKey extends Array<unknown> = unknown[]>({ dagIds }: {
737+
dagIds: string[];
738+
}, queryKey?: TQueryKey, options?: Omit<UseQueryOptions<TData, TError>, "queryKey" | "queryFn">) => useSuspenseQuery<TData, TError>({ queryKey: Common.UseDagServiceGetLatestRunTaskInstanceStateCountsUiKeyFn({ dagIds }, queryKey), queryFn: () => DagService.getLatestRunTaskInstanceStateCountsUi({ dagIds }) as TData, ...options });
739+
/**
727740
* Get Event Log
728741
* @param data The data for the request.
729742
* @param data.eventLogId

‎airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8865,6 +8865,33 @@ It is used to transfer providers information loaded by providers_manager such th
88658865
the API server/Web UI can use this data to render connection form UI.`
88668866
} as const;
88678867

8868+
export const $DAGLatestRunTaskInstanceStateCountsResponse = {
8869+
properties: {
8870+
dag_id: {
8871+
type: 'string',
8872+
title: 'Dag Id'
8873+
},
8874+
run_id: {
8875+
type: 'string',
8876+
title: 'Run Id'
8877+
},
8878+
state_counts: {
8879+
additionalProperties: {
8880+
type: 'integer'
8881+
},
8882+
type: 'object',
8883+
title: 'State Counts'
8884+
}
8885+
},
8886+
type: 'object',
8887+
required: ['dag_id', 'run_id', 'state_counts'],
8888+
title: 'DAGLatestRunTaskInstanceStateCountsResponse',
8889+
description: `Task-instance state counts for a Dag's latest run.
8890+
8891+
\`\`state_counts\`\` only carries states present in the run; task instances without a
8892+
state yet are keyed as \`\`no_status\`\`.`
8893+
} as const;
8894+
88688895
export const $DAGRunLightResponse = {
88698896
properties: {
88708897
id: {
@@ -9315,6 +9342,22 @@ export const $DAGWithLatestDagRunsResponse = {
93159342
description: 'DAG with latest dag runs response serializer.'
93169343
} as const;
93179344

9345+
export const $DAGsLatestRunTaskInstanceStateCountsCollectionResponse = {
9346+
properties: {
9347+
dags: {
9348+
items: {
9349+
'$ref': '#/components/schemas/DAGLatestRunTaskInstanceStateCountsResponse'
9350+
},
9351+
type: 'array',
9352+
title: 'Dags'
9353+
}
9354+
},
9355+
type: 'object',
9356+
required: ['dags'],
9357+
title: 'DAGsLatestRunTaskInstanceStateCountsCollectionResponse',
9358+
description: 'Collection of per-Dag latest-run task-instance state counts for the Dag list page.'
9359+
} as const;
9360+
93189361
export const $DAGsRunStateCountsCollectionResponse = {
93199362
properties: {
93209363
dags: {

0 commit comments

Comments
 (0)