Skip to content

Commit 4ed8bdd

Browse files
pierrejeambrunpierrejeambrun
authored
Update UI list deadlines endpoint (#62374)
* Adjustments to ui list deadlines endpoint * Update tests
1 parent 6374e4d commit 4ed8bdd

File tree

12 files changed

+264
-100
lines changed

12 files changed

+264
-100
lines changed

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

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,18 +17,30 @@
1717

1818
from __future__ import annotations
1919

20+
from collections.abc import Iterable
2021
from datetime import datetime
2122
from uuid import UUID
2223

24+
from pydantic import AliasPath, Field
25+
2326
from airflow.api_fastapi.core_api.base import BaseModel
2427

2528

2629
class DeadlineResponse(BaseModel):
27-
"""Deadline data for the DAG run deadlines tab."""
30+
"""Deadline serializer for responses."""
2831

2932
id: UUID
3033
deadline_time: datetime
3134
missed: bool
3235
created_at: datetime
33-
alert_name: str | None = None
34-
alert_description: str | None = None
36+
alert_name: str | None = Field(validation_alias=AliasPath("deadline_alert", "name"), default=None)
37+
alert_description: str | None = Field(
38+
validation_alias=AliasPath("deadline_alert", "description"), default=None
39+
)
40+
41+
42+
class DealineCollectionResponse(BaseModel):
43+
"""Deadline Collection serializer for responses."""
44+
45+
deadlines: Iterable[DeadlineResponse]
46+
total_entries: int

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

Lines changed: 53 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -554,7 +554,7 @@ paths:
554554
security:
555555
- OAuth2PasswordBearer: []
556556
- HTTPBearer: []
557-
/ui/deadlines/{dag_id}/{run_id}:
557+
/ui/dags/{dag_id}/dagRuns/{dag_run_id}/deadlines:
558558
get:
559559
tags:
560560
- Deadlines
@@ -571,22 +571,51 @@ paths:
571571
schema:
572572
type: string
573573
title: Dag Id
574-
- name: run_id
574+
- name: dag_run_id
575575
in: path
576576
required: true
577577
schema:
578578
type: string
579-
title: Run Id
579+
title: Dag Run Id
580+
- name: limit
581+
in: query
582+
required: false
583+
schema:
584+
type: integer
585+
minimum: 0
586+
default: 50
587+
title: Limit
588+
- name: offset
589+
in: query
590+
required: false
591+
schema:
592+
type: integer
593+
minimum: 0
594+
default: 0
595+
title: Offset
596+
- name: order_by
597+
in: query
598+
required: false
599+
schema:
600+
type: array
601+
items:
602+
type: string
603+
description: 'Attributes to order by, multi criteria sort is supported.
604+
Prefix with `-` for descending order. Supported attributes: `id, deadline_time,
605+
created_at, alert_name`'
606+
default:
607+
- deadline_time
608+
title: Order By
609+
description: 'Attributes to order by, multi criteria sort is supported. Prefix
610+
with `-` for descending order. Supported attributes: `id, deadline_time,
611+
created_at, alert_name`'
580612
responses:
581613
'200':
582614
description: Successful Response
583615
content:
584616
application/json:
585617
schema:
586-
type: array
587-
items:
588-
$ref: '#/components/schemas/DeadlineResponse'
589-
title: Response Get Dag Run Deadlines
618+
$ref: '#/components/schemas/DealineCollectionResponse'
590619
'404':
591620
content:
592621
application/json:
@@ -2081,7 +2110,23 @@ components:
20812110
- missed
20822111
- created_at
20832112
title: DeadlineResponse
2084-
description: Deadline data for the DAG run deadlines tab.
2113+
description: Deadline serializer for responses.
2114+
DealineCollectionResponse:
2115+
properties:
2116+
deadlines:
2117+
items:
2118+
$ref: '#/components/schemas/DeadlineResponse'
2119+
type: array
2120+
title: Deadlines
2121+
total_entries:
2122+
type: integer
2123+
title: Total Entries
2124+
type: object
2125+
required:
2126+
- deadlines
2127+
- total_entries
2128+
title: DealineCollectionResponse
2129+
description: Deadline Collection serializer for responses.
20852130
EdgeResponse:
20862131
properties:
20872132
source_id:

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

Lines changed: 46 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -17,24 +17,28 @@
1717

1818
from __future__ import annotations
1919

20+
from typing import Annotated
21+
2022
from fastapi import Depends, HTTPException, status
2123
from sqlalchemy import select
2224
from sqlalchemy.orm import joinedload
2325

2426
from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity
25-
from airflow.api_fastapi.common.db.common import SessionDep
27+
from airflow.api_fastapi.common.db.common import SessionDep, paginated_select
28+
from airflow.api_fastapi.common.parameters import QueryLimit, QueryOffset, SortParam
2629
from airflow.api_fastapi.common.router import AirflowRouter
27-
from airflow.api_fastapi.core_api.datamodels.ui.deadline import DeadlineResponse
30+
from airflow.api_fastapi.core_api.datamodels.ui.deadline import DealineCollectionResponse
2831
from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
2932
from airflow.api_fastapi.core_api.security import requires_access_dag
3033
from airflow.models.dagrun import DagRun
3134
from airflow.models.deadline import Deadline
35+
from airflow.models.deadline_alert import DeadlineAlert
3236

33-
deadlines_router = AirflowRouter(prefix="/deadlines", tags=["Deadlines"])
37+
deadlines_router = AirflowRouter(prefix="/dags/{dag_id}/dagRuns/{dag_run_id}/deadlines", tags=["Deadlines"])
3438

3539

3640
@deadlines_router.get(
37-
"/{dag_id}/{run_id}",
41+
"",
3842
responses=create_openapi_http_exception_doc(
3943
[
4044
status.HTTP_404_NOT_FOUND,
@@ -51,36 +55,50 @@
5155
)
5256
def get_dag_run_deadlines(
5357
dag_id: str,
54-
run_id: str,
58+
dag_run_id: str,
5559
session: SessionDep,
56-
) -> list[DeadlineResponse]:
60+
limit: QueryLimit,
61+
offset: QueryOffset,
62+
order_by: Annotated[
63+
SortParam,
64+
Depends(
65+
SortParam(
66+
["id", "deadline_time", "created_at"],
67+
Deadline,
68+
to_replace={
69+
"alert_name": DeadlineAlert.name,
70+
},
71+
).dynamic_depends(default="deadline_time")
72+
),
73+
],
74+
) -> DealineCollectionResponse:
5775
"""Get all deadlines for a specific DAG run."""
58-
dag_run = session.scalar(select(DagRun).where(DagRun.dag_id == dag_id, DagRun.run_id == run_id))
76+
dag_run = session.scalar(select(DagRun).where(DagRun.dag_id == dag_id, DagRun.run_id == dag_run_id))
77+
5978
if not dag_run:
6079
raise HTTPException(
6180
status.HTTP_404_NOT_FOUND,
62-
f"No DAG run found for dag_id={dag_id} run_id={run_id}",
81+
f"No DAG run found for dag_id={dag_id} dag_run_id={dag_run_id}",
6382
)
6483

65-
deadlines = (
66-
session.scalars(
67-
select(Deadline)
68-
.where(Deadline.dagrun_id == dag_run.id)
69-
.options(joinedload(Deadline.deadline_alert))
70-
.order_by(Deadline.deadline_time.asc())
71-
)
72-
.unique()
73-
.all()
84+
query = (
85+
select(Deadline)
86+
.join(Deadline.dagrun)
87+
.outerjoin(Deadline.deadline_alert)
88+
.where(Deadline.dagrun_id == dag_run.id)
89+
.where(DagRun.dag_id == dag_id)
90+
.options(joinedload(Deadline.deadline_alert))
7491
)
7592

76-
return [
77-
DeadlineResponse(
78-
id=d.id,
79-
deadline_time=d.deadline_time,
80-
missed=d.missed,
81-
created_at=d.created_at,
82-
alert_name=d.deadline_alert.name if d.deadline_alert else None,
83-
alert_description=d.deadline_alert.description if d.deadline_alert else None,
84-
)
85-
for d in deadlines
86-
]
93+
deadlines_select, total_entries = paginated_select(
94+
statement=query,
95+
filters=None,
96+
order_by=order_by,
97+
offset=offset,
98+
limit=limit,
99+
session=session,
100+
)
101+
102+
deadlines = session.scalars(deadlines_select)
103+
104+
return DealineCollectionResponse(deadlines=deadlines, total_entries=total_entries)

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

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -824,10 +824,13 @@ export const UseDashboardServiceDagStatsKeyFn = (queryKey?: Array<unknown>) => [
824824
export type DeadlinesServiceGetDagRunDeadlinesDefaultResponse = Awaited<ReturnType<typeof DeadlinesService.getDagRunDeadlines>>;
825825
export type DeadlinesServiceGetDagRunDeadlinesQueryResult<TData = DeadlinesServiceGetDagRunDeadlinesDefaultResponse, TError = unknown> = UseQueryResult<TData, TError>;
826826
export const useDeadlinesServiceGetDagRunDeadlinesKey = "DeadlinesServiceGetDagRunDeadlines";
827-
export const UseDeadlinesServiceGetDagRunDeadlinesKeyFn = ({ dagId, runId }: {
827+
export const UseDeadlinesServiceGetDagRunDeadlinesKeyFn = ({ dagId, dagRunId, limit, offset, orderBy }: {
828828
dagId: string;
829-
runId: string;
830-
}, queryKey?: Array<unknown>) => [useDeadlinesServiceGetDagRunDeadlinesKey, ...(queryKey ?? [{ dagId, runId }])];
829+
dagRunId: string;
830+
limit?: number;
831+
offset?: number;
832+
orderBy?: string[];
833+
}, queryKey?: Array<unknown>) => [useDeadlinesServiceGetDagRunDeadlinesKey, ...(queryKey ?? [{ dagId, dagRunId, limit, offset, orderBy }])];
831834
export type StructureServiceStructureDataDefaultResponse = Awaited<ReturnType<typeof StructureService.structureData>>;
832835
export type StructureServiceStructureDataQueryResult<TData = StructureServiceStructureDataDefaultResponse, TError = unknown> = UseQueryResult<TData, TError>;
833836
export const useStructureServiceStructureDataKey = "StructureServiceStructureData";

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

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1562,14 +1562,20 @@ export const ensureUseDashboardServiceDagStatsData = (queryClient: QueryClient)
15621562
* Get all deadlines for a specific DAG run.
15631563
* @param data The data for the request.
15641564
* @param data.dagId
1565-
* @param data.runId
1566-
* @returns DeadlineResponse Successful Response
1565+
* @param data.dagRunId
1566+
* @param data.limit
1567+
* @param data.offset
1568+
* @param data.orderBy Attributes to order by, multi criteria sort is supported. Prefix with `-` for descending order. Supported attributes: `id, deadline_time, created_at, alert_name`
1569+
* @returns DealineCollectionResponse Successful Response
15671570
* @throws ApiError
15681571
*/
1569-
export const ensureUseDeadlinesServiceGetDagRunDeadlinesData = (queryClient: QueryClient, { dagId, runId }: {
1572+
export const ensureUseDeadlinesServiceGetDagRunDeadlinesData = (queryClient: QueryClient, { dagId, dagRunId, limit, offset, orderBy }: {
15701573
dagId: string;
1571-
runId: string;
1572-
}) => queryClient.ensureQueryData({ queryKey: Common.UseDeadlinesServiceGetDagRunDeadlinesKeyFn({ dagId, runId }), queryFn: () => DeadlinesService.getDagRunDeadlines({ dagId, runId }) });
1574+
dagRunId: string;
1575+
limit?: number;
1576+
offset?: number;
1577+
orderBy?: string[];
1578+
}) => queryClient.ensureQueryData({ queryKey: Common.UseDeadlinesServiceGetDagRunDeadlinesKeyFn({ dagId, dagRunId, limit, offset, orderBy }), queryFn: () => DeadlinesService.getDagRunDeadlines({ dagId, dagRunId, limit, offset, orderBy }) });
15731579
/**
15741580
* Structure Data
15751581
* Get Structure Data.

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

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1562,14 +1562,20 @@ export const prefetchUseDashboardServiceDagStats = (queryClient: QueryClient) =>
15621562
* Get all deadlines for a specific DAG run.
15631563
* @param data The data for the request.
15641564
* @param data.dagId
1565-
* @param data.runId
1566-
* @returns DeadlineResponse Successful Response
1565+
* @param data.dagRunId
1566+
* @param data.limit
1567+
* @param data.offset
1568+
* @param data.orderBy Attributes to order by, multi criteria sort is supported. Prefix with `-` for descending order. Supported attributes: `id, deadline_time, created_at, alert_name`
1569+
* @returns DealineCollectionResponse Successful Response
15671570
* @throws ApiError
15681571
*/
1569-
export const prefetchUseDeadlinesServiceGetDagRunDeadlines = (queryClient: QueryClient, { dagId, runId }: {
1572+
export const prefetchUseDeadlinesServiceGetDagRunDeadlines = (queryClient: QueryClient, { dagId, dagRunId, limit, offset, orderBy }: {
15701573
dagId: string;
1571-
runId: string;
1572-
}) => queryClient.prefetchQuery({ queryKey: Common.UseDeadlinesServiceGetDagRunDeadlinesKeyFn({ dagId, runId }), queryFn: () => DeadlinesService.getDagRunDeadlines({ dagId, runId }) });
1574+
dagRunId: string;
1575+
limit?: number;
1576+
offset?: number;
1577+
orderBy?: string[];
1578+
}) => queryClient.prefetchQuery({ queryKey: Common.UseDeadlinesServiceGetDagRunDeadlinesKeyFn({ dagId, dagRunId, limit, offset, orderBy }), queryFn: () => DeadlinesService.getDagRunDeadlines({ dagId, dagRunId, limit, offset, orderBy }) });
15731579
/**
15741580
* Structure Data
15751581
* Get Structure Data.

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

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1562,14 +1562,20 @@ export const useDashboardServiceDagStats = <TData = Common.DashboardServiceDagSt
15621562
* Get all deadlines for a specific DAG run.
15631563
* @param data The data for the request.
15641564
* @param data.dagId
1565-
* @param data.runId
1566-
* @returns DeadlineResponse Successful Response
1565+
* @param data.dagRunId
1566+
* @param data.limit
1567+
* @param data.offset
1568+
* @param data.orderBy Attributes to order by, multi criteria sort is supported. Prefix with `-` for descending order. Supported attributes: `id, deadline_time, created_at, alert_name`
1569+
* @returns DealineCollectionResponse Successful Response
15671570
* @throws ApiError
15681571
*/
1569-
export const useDeadlinesServiceGetDagRunDeadlines = <TData = Common.DeadlinesServiceGetDagRunDeadlinesDefaultResponse, TError = unknown, TQueryKey extends Array<unknown> = unknown[]>({ dagId, runId }: {
1572+
export const useDeadlinesServiceGetDagRunDeadlines = <TData = Common.DeadlinesServiceGetDagRunDeadlinesDefaultResponse, TError = unknown, TQueryKey extends Array<unknown> = unknown[]>({ dagId, dagRunId, limit, offset, orderBy }: {
15701573
dagId: string;
1571-
runId: string;
1572-
}, queryKey?: TQueryKey, options?: Omit<UseQueryOptions<TData, TError>, "queryKey" | "queryFn">) => useQuery<TData, TError>({ queryKey: Common.UseDeadlinesServiceGetDagRunDeadlinesKeyFn({ dagId, runId }, queryKey), queryFn: () => DeadlinesService.getDagRunDeadlines({ dagId, runId }) as TData, ...options });
1574+
dagRunId: string;
1575+
limit?: number;
1576+
offset?: number;
1577+
orderBy?: string[];
1578+
}, queryKey?: TQueryKey, options?: Omit<UseQueryOptions<TData, TError>, "queryKey" | "queryFn">) => useQuery<TData, TError>({ queryKey: Common.UseDeadlinesServiceGetDagRunDeadlinesKeyFn({ dagId, dagRunId, limit, offset, orderBy }, queryKey), queryFn: () => DeadlinesService.getDagRunDeadlines({ dagId, dagRunId, limit, offset, orderBy }) as TData, ...options });
15731579
/**
15741580
* Structure Data
15751581
* Get Structure Data.

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

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1562,14 +1562,20 @@ export const useDashboardServiceDagStatsSuspense = <TData = Common.DashboardServ
15621562
* Get all deadlines for a specific DAG run.
15631563
* @param data The data for the request.
15641564
* @param data.dagId
1565-
* @param data.runId
1566-
* @returns DeadlineResponse Successful Response
1565+
* @param data.dagRunId
1566+
* @param data.limit
1567+
* @param data.offset
1568+
* @param data.orderBy Attributes to order by, multi criteria sort is supported. Prefix with `-` for descending order. Supported attributes: `id, deadline_time, created_at, alert_name`
1569+
* @returns DealineCollectionResponse Successful Response
15671570
* @throws ApiError
15681571
*/
1569-
export const useDeadlinesServiceGetDagRunDeadlinesSuspense = <TData = Common.DeadlinesServiceGetDagRunDeadlinesDefaultResponse, TError = unknown, TQueryKey extends Array<unknown> = unknown[]>({ dagId, runId }: {
1572+
export const useDeadlinesServiceGetDagRunDeadlinesSuspense = <TData = Common.DeadlinesServiceGetDagRunDeadlinesDefaultResponse, TError = unknown, TQueryKey extends Array<unknown> = unknown[]>({ dagId, dagRunId, limit, offset, orderBy }: {
15701573
dagId: string;
1571-
runId: string;
1572-
}, queryKey?: TQueryKey, options?: Omit<UseQueryOptions<TData, TError>, "queryKey" | "queryFn">) => useSuspenseQuery<TData, TError>({ queryKey: Common.UseDeadlinesServiceGetDagRunDeadlinesKeyFn({ dagId, runId }, queryKey), queryFn: () => DeadlinesService.getDagRunDeadlines({ dagId, runId }) as TData, ...options });
1574+
dagRunId: string;
1575+
limit?: number;
1576+
offset?: number;
1577+
orderBy?: string[];
1578+
}, queryKey?: TQueryKey, options?: Omit<UseQueryOptions<TData, TError>, "queryKey" | "queryFn">) => useSuspenseQuery<TData, TError>({ queryKey: Common.UseDeadlinesServiceGetDagRunDeadlinesKeyFn({ dagId, dagRunId, limit, offset, orderBy }, queryKey), queryFn: () => DeadlinesService.getDagRunDeadlines({ dagId, dagRunId, limit, offset, orderBy }) as TData, ...options });
15731579
/**
15741580
* Structure Data
15751581
* Get Structure Data.

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

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7897,7 +7897,27 @@ export const $DeadlineResponse = {
78977897
type: 'object',
78987898
required: ['id', 'deadline_time', 'missed', 'created_at'],
78997899
title: 'DeadlineResponse',
7900-
description: 'Deadline data for the DAG run deadlines tab.'
7900+
description: 'Deadline serializer for responses.'
7901+
} as const;
7902+
7903+
export const $DealineCollectionResponse = {
7904+
properties: {
7905+
deadlines: {
7906+
items: {
7907+
'$ref': '#/components/schemas/DeadlineResponse'
7908+
},
7909+
type: 'array',
7910+
title: 'Deadlines'
7911+
},
7912+
total_entries: {
7913+
type: 'integer',
7914+
title: 'Total Entries'
7915+
}
7916+
},
7917+
type: 'object',
7918+
required: ['deadlines', 'total_entries'],
7919+
title: 'DealineCollectionResponse',
7920+
description: 'Deadline Collection serializer for responses.'
79017921
} as const;
79027922

79037923
export const $EdgeResponse = {

0 commit comments

Comments
 (0)