This is an automated email from the ASF dual-hosted git repository.
bbovenzi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new f806187dc85 UI: Add team column and filter to the deadlines list
(#72060)
f806187dc85 is described below
commit f806187dc85343d82fa4ab477ae5e7c71c1c50f1
Author: Vincent <[email protected]>
AuthorDate: Mon Aug 31 16:25:01 2026 -0400
UI: Add team column and filter to the deadlines list (#72060)
* UI: Add team column and filter to the deadlines list
When multi-team mode is enabled, operators watching deadlines across teams
need
to see which team owns each deadline and to scope the list to a single team
when
triaging that team's misses.
The column and the filter render only when the `multi_team` configuration is
enabled, and no team data is loaded when it is off, so single-team
deployments
are unaffected.
* Make deadlines UI tests recover from leftover team rows
attach_dag_to_team commits its bundle and team before yielding, so a test
run
interrupted inside that context manager leaves the rows behind. On a
persistent
database every later run then failed with a duplicate-key error until the
rows
were deleted by hand.
---
.../api_fastapi/core_api/datamodels/ui/deadline.py | 1 +
.../api_fastapi/core_api/openapi/_private_ui.yaml | 13 ++++
.../api_fastapi/core_api/routes/ui/deadlines.py | 10 ++-
.../src/airflow/ui/openapi-gen/queries/common.ts | 5 +-
.../ui/openapi-gen/queries/ensureQueryData.ts | 6 +-
.../src/airflow/ui/openapi-gen/queries/prefetch.ts | 6 +-
.../src/airflow/ui/openapi-gen/queries/queries.ts | 6 +-
.../src/airflow/ui/openapi-gen/queries/suspense.ts | 6 +-
.../airflow/ui/openapi-gen/requests/schemas.gen.ts | 11 +++
.../ui/openapi-gen/requests/services.gen.ts | 4 +-
.../airflow/ui/openapi-gen/requests/types.gen.ts | 2 +
.../src/airflow/ui/src/pages/Deadlines/index.tsx | 23 +++++-
.../core_api/routes/public/test_dag_run.py | 71 ++-----------------
.../core_api/routes/public/test_task_instances.py | 82 +++-------------------
.../core_api/routes/ui/test_deadlines.py | 63 +++++++++++++++--
devel-common/src/tests_common/test_utils/team.py | 68 ++++++++++++++++++
16 files changed, 215 insertions(+), 162 deletions(-)
diff --git
a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/deadline.py
b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/deadline.py
index 9c09176fea4..53b583d1f40 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/deadline.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/deadline.py
@@ -38,6 +38,7 @@ class DeadlineResponse(BaseModel):
dag_run_id: str = Field(validation_alias=AliasPath("dagrun", "run_id"))
alert_id: UUID | None = Field(validation_alias="deadline_alert_id",
default=None)
alert_name: str | None =
Field(validation_alias=AliasPath("deadline_alert", "name"), default=None)
+ team_name: str | None = Field(validation_alias=AliasPath("dagrun",
"team_name"), default=None)
class DeadlineCollectionResponse(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 25b50d9e825..6ca0e3c24b4 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
@@ -1142,6 +1142,14 @@ paths:
format: date-time
- type: 'null'
title: Last Updated At Lt
+ - name: teams
+ in: query
+ required: false
+ schema:
+ type: array
+ items:
+ type: string
+ title: Teams
responses:
'200':
description: Successful Response
@@ -3368,6 +3376,11 @@ components:
- type: string
- type: 'null'
title: Alert Name
+ team_name:
+ anyOf:
+ - type: string
+ - type: 'null'
+ title: Team Name
type: object
required:
- id
diff --git
a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/deadlines.py
b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/deadlines.py
index 4127e98104e..6a239540a75 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/deadlines.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/deadlines.py
@@ -25,14 +25,17 @@ from sqlalchemy.orm import contains_eager, noload
from airflow.api_fastapi.auth.managers.models.resource_details import
DagAccessEntity
from airflow.api_fastapi.common.db.common import SessionDep, paginated_select
+from airflow.api_fastapi.common.db.dags import eager_load_teams
from airflow.api_fastapi.common.parameters import (
FilterParam,
QueryLimit,
QueryOffset,
RangeFilter,
SortParam,
+ _DagIdTeamsFilter,
datetime_range_filter_factory,
filter_param_factory,
+ teams_filter_factory,
)
from airflow.api_fastapi.common.router import AirflowRouter
from airflow.api_fastapi.core_api.datamodels.ui.deadline import (
@@ -93,6 +96,7 @@ def get_deadlines(
last_updated_at: Annotated[
RangeFilter, Depends(datetime_range_filter_factory("last_updated_at",
Deadline))
],
+ teams: Annotated[_DagIdTeamsFilter,
Depends(teams_filter_factory(DagRun.dag_id))],
) -> DeadlineCollectionResponse:
"""
Get deadlines for a Dag run.
@@ -105,7 +109,9 @@ def get_deadlines(
.join(Deadline.dagrun)
.outerjoin(Deadline.deadline_alert)
.options(
- contains_eager(Deadline.dagrun).options(noload(DagRun.deadlines)),
+ contains_eager(Deadline.dagrun).options(
+ noload(DagRun.deadlines), *eager_load_teams(DagRun.dag_model)
+ ),
contains_eager(Deadline.deadline_alert),
noload(Deadline.callback),
)
@@ -123,7 +129,7 @@ def get_deadlines(
deadlines_select, total_entries = paginated_select(
statement=query,
- filters=[readable_dag_runs_filter, missed, deadline_time,
last_updated_at],
+ filters=[readable_dag_runs_filter, missed, deadline_time,
last_updated_at, teams],
order_by=order_by,
offset=offset,
limit=limit,
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts
b/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts
index e5c078e4a11..206e6c29c87 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts
@@ -986,7 +986,7 @@ export const UseDashboardServiceDagStatsKeyFn = (queryKey?:
Array<unknown>) => [
export type DeadlinesServiceGetDeadlinesDefaultResponse =
Awaited<ReturnType<typeof DeadlinesService.getDeadlines>>;
export type DeadlinesServiceGetDeadlinesQueryResult<TData =
DeadlinesServiceGetDeadlinesDefaultResponse, TError = unknown> =
UseQueryResult<TData, TError>;
export const useDeadlinesServiceGetDeadlinesKey =
"DeadlinesServiceGetDeadlines";
-export const UseDeadlinesServiceGetDeadlinesKeyFn = ({ dagId, dagRunId,
deadlineTimeGt, deadlineTimeGte, deadlineTimeLt, deadlineTimeLte,
lastUpdatedAtGt, lastUpdatedAtGte, lastUpdatedAtLt, lastUpdatedAtLte, limit,
missed, offset, orderBy }: {
+export const UseDeadlinesServiceGetDeadlinesKeyFn = ({ dagId, dagRunId,
deadlineTimeGt, deadlineTimeGte, deadlineTimeLt, deadlineTimeLte,
lastUpdatedAtGt, lastUpdatedAtGte, lastUpdatedAtLt, lastUpdatedAtLte, limit,
missed, offset, orderBy, teams }: {
dagId: string;
dagRunId: string;
deadlineTimeGt?: string;
@@ -1001,7 +1001,8 @@ export const UseDeadlinesServiceGetDeadlinesKeyFn = ({
dagId, dagRunId, deadline
missed?: boolean;
offset?: number;
orderBy?: string[];
-}, queryKey?: Array<unknown>) => [useDeadlinesServiceGetDeadlinesKey,
...(queryKey ?? [{ dagId, dagRunId, deadlineTimeGt, deadlineTimeGte,
deadlineTimeLt, deadlineTimeLte, lastUpdatedAtGt, lastUpdatedAtGte,
lastUpdatedAtLt, lastUpdatedAtLte, limit, missed, offset, orderBy }])];
+ teams?: string[];
+}, queryKey?: Array<unknown>) => [useDeadlinesServiceGetDeadlinesKey,
...(queryKey ?? [{ dagId, dagRunId, deadlineTimeGt, deadlineTimeGte,
deadlineTimeLt, deadlineTimeLte, lastUpdatedAtGt, lastUpdatedAtGte,
lastUpdatedAtLt, lastUpdatedAtLte, limit, missed, offset, orderBy, teams }])];
export type DeadlinesServiceGetDagDeadlineAlertsDefaultResponse =
Awaited<ReturnType<typeof DeadlinesService.getDagDeadlineAlerts>>;
export type DeadlinesServiceGetDagDeadlineAlertsQueryResult<TData =
DeadlinesServiceGetDagDeadlineAlertsDefaultResponse, TError = unknown> =
UseQueryResult<TData, TError>;
export const useDeadlinesServiceGetDagDeadlineAlertsKey =
"DeadlinesServiceGetDagDeadlineAlerts";
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
index 2da84b8596d..bd88502f1e1 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
@@ -1912,10 +1912,11 @@ export const ensureUseDashboardServiceDagStatsData =
(queryClient: QueryClient)
* @param data.lastUpdatedAtGt
* @param data.lastUpdatedAtLte
* @param data.lastUpdatedAtLt
+* @param data.teams
* @returns DeadlineCollectionResponse Successful Response
* @throws ApiError
*/
-export const ensureUseDeadlinesServiceGetDeadlinesData = (queryClient:
QueryClient, { dagId, dagRunId, deadlineTimeGt, deadlineTimeGte,
deadlineTimeLt, deadlineTimeLte, lastUpdatedAtGt, lastUpdatedAtGte,
lastUpdatedAtLt, lastUpdatedAtLte, limit, missed, offset, orderBy }: {
+export const ensureUseDeadlinesServiceGetDeadlinesData = (queryClient:
QueryClient, { dagId, dagRunId, deadlineTimeGt, deadlineTimeGte,
deadlineTimeLt, deadlineTimeLte, lastUpdatedAtGt, lastUpdatedAtGte,
lastUpdatedAtLt, lastUpdatedAtLte, limit, missed, offset, orderBy, teams }: {
dagId: string;
dagRunId: string;
deadlineTimeGt?: string;
@@ -1930,7 +1931,8 @@ export const ensureUseDeadlinesServiceGetDeadlinesData =
(queryClient: QueryClie
missed?: boolean;
offset?: number;
orderBy?: string[];
-}) => queryClient.ensureQueryData({ queryKey:
Common.UseDeadlinesServiceGetDeadlinesKeyFn({ dagId, dagRunId, deadlineTimeGt,
deadlineTimeGte, deadlineTimeLt, deadlineTimeLte, lastUpdatedAtGt,
lastUpdatedAtGte, lastUpdatedAtLt, lastUpdatedAtLte, limit, missed, offset,
orderBy }), queryFn: () => DeadlinesService.getDeadlines({ dagId, dagRunId,
deadlineTimeGt, deadlineTimeGte, deadlineTimeLt, deadlineTimeLte,
lastUpdatedAtGt, lastUpdatedAtGte, lastUpdatedAtLt, lastUpdatedAtLte, limit,
misse [...]
+ teams?: string[];
+}) => queryClient.ensureQueryData({ queryKey:
Common.UseDeadlinesServiceGetDeadlinesKeyFn({ dagId, dagRunId, deadlineTimeGt,
deadlineTimeGte, deadlineTimeLt, deadlineTimeLte, lastUpdatedAtGt,
lastUpdatedAtGte, lastUpdatedAtLt, lastUpdatedAtLte, limit, missed, offset,
orderBy, teams }), queryFn: () => DeadlinesService.getDeadlines({ dagId,
dagRunId, deadlineTimeGt, deadlineTimeGte, deadlineTimeLt, deadlineTimeLte,
lastUpdatedAtGt, lastUpdatedAtGte, lastUpdatedAtLt, lastUpdatedAtLte, limit
[...]
/**
* Get Dag Deadline Alerts
* Get all deadline alerts defined on a Dag.
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
index 917cc830f04..8d28a3aa001 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
@@ -1912,10 +1912,11 @@ export const prefetchUseDashboardServiceDagStats =
(queryClient: QueryClient) =>
* @param data.lastUpdatedAtGt
* @param data.lastUpdatedAtLte
* @param data.lastUpdatedAtLt
+* @param data.teams
* @returns DeadlineCollectionResponse Successful Response
* @throws ApiError
*/
-export const prefetchUseDeadlinesServiceGetDeadlines = (queryClient:
QueryClient, { dagId, dagRunId, deadlineTimeGt, deadlineTimeGte,
deadlineTimeLt, deadlineTimeLte, lastUpdatedAtGt, lastUpdatedAtGte,
lastUpdatedAtLt, lastUpdatedAtLte, limit, missed, offset, orderBy }: {
+export const prefetchUseDeadlinesServiceGetDeadlines = (queryClient:
QueryClient, { dagId, dagRunId, deadlineTimeGt, deadlineTimeGte,
deadlineTimeLt, deadlineTimeLte, lastUpdatedAtGt, lastUpdatedAtGte,
lastUpdatedAtLt, lastUpdatedAtLte, limit, missed, offset, orderBy, teams }: {
dagId: string;
dagRunId: string;
deadlineTimeGt?: string;
@@ -1930,7 +1931,8 @@ export const prefetchUseDeadlinesServiceGetDeadlines =
(queryClient: QueryClient
missed?: boolean;
offset?: number;
orderBy?: string[];
-}) => queryClient.prefetchQuery({ queryKey:
Common.UseDeadlinesServiceGetDeadlinesKeyFn({ dagId, dagRunId, deadlineTimeGt,
deadlineTimeGte, deadlineTimeLt, deadlineTimeLte, lastUpdatedAtGt,
lastUpdatedAtGte, lastUpdatedAtLt, lastUpdatedAtLte, limit, missed, offset,
orderBy }), queryFn: () => DeadlinesService.getDeadlines({ dagId, dagRunId,
deadlineTimeGt, deadlineTimeGte, deadlineTimeLt, deadlineTimeLte,
lastUpdatedAtGt, lastUpdatedAtGte, lastUpdatedAtLt, lastUpdatedAtLte, limit,
missed, [...]
+ teams?: string[];
+}) => queryClient.prefetchQuery({ queryKey:
Common.UseDeadlinesServiceGetDeadlinesKeyFn({ dagId, dagRunId, deadlineTimeGt,
deadlineTimeGte, deadlineTimeLt, deadlineTimeLte, lastUpdatedAtGt,
lastUpdatedAtGte, lastUpdatedAtLt, lastUpdatedAtLte, limit, missed, offset,
orderBy, teams }), queryFn: () => DeadlinesService.getDeadlines({ dagId,
dagRunId, deadlineTimeGt, deadlineTimeGte, deadlineTimeLt, deadlineTimeLte,
lastUpdatedAtGt, lastUpdatedAtGte, lastUpdatedAtLt, lastUpdatedAtLte, limit,
[...]
/**
* Get Dag Deadline Alerts
* Get all deadline alerts defined on a Dag.
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
index 3ea1aede3b3..9c7971fb2d6 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
@@ -1912,10 +1912,11 @@ export const useDashboardServiceDagStats = <TData =
Common.DashboardServiceDagSt
* @param data.lastUpdatedAtGt
* @param data.lastUpdatedAtLte
* @param data.lastUpdatedAtLt
+* @param data.teams
* @returns DeadlineCollectionResponse Successful Response
* @throws ApiError
*/
-export const useDeadlinesServiceGetDeadlines = <TData =
Common.DeadlinesServiceGetDeadlinesDefaultResponse, TError = unknown, TQueryKey
extends Array<unknown> = unknown[]>({ dagId, dagRunId, deadlineTimeGt,
deadlineTimeGte, deadlineTimeLt, deadlineTimeLte, lastUpdatedAtGt,
lastUpdatedAtGte, lastUpdatedAtLt, lastUpdatedAtLte, limit, missed, offset,
orderBy }: {
+export const useDeadlinesServiceGetDeadlines = <TData =
Common.DeadlinesServiceGetDeadlinesDefaultResponse, TError = unknown, TQueryKey
extends Array<unknown> = unknown[]>({ dagId, dagRunId, deadlineTimeGt,
deadlineTimeGte, deadlineTimeLt, deadlineTimeLte, lastUpdatedAtGt,
lastUpdatedAtGte, lastUpdatedAtLt, lastUpdatedAtLte, limit, missed, offset,
orderBy, teams }: {
dagId: string;
dagRunId: string;
deadlineTimeGt?: string;
@@ -1930,7 +1931,8 @@ export const useDeadlinesServiceGetDeadlines = <TData =
Common.DeadlinesServiceG
missed?: boolean;
offset?: number;
orderBy?: string[];
-}, queryKey?: TQueryKey, options?: Omit<UseQueryOptions<TData, TError>,
"queryKey" | "queryFn">) => useQuery<TData, TError>({ queryKey:
Common.UseDeadlinesServiceGetDeadlinesKeyFn({ dagId, dagRunId, deadlineTimeGt,
deadlineTimeGte, deadlineTimeLt, deadlineTimeLte, lastUpdatedAtGt,
lastUpdatedAtGte, lastUpdatedAtLt, lastUpdatedAtLte, limit, missed, offset,
orderBy }, queryKey), queryFn: () => DeadlinesService.getDeadlines({ dagId,
dagRunId, deadlineTimeGt, deadlineTimeGte, deadlineTimeLt, [...]
+ teams?: string[];
+}, queryKey?: TQueryKey, options?: Omit<UseQueryOptions<TData, TError>,
"queryKey" | "queryFn">) => useQuery<TData, TError>({ queryKey:
Common.UseDeadlinesServiceGetDeadlinesKeyFn({ dagId, dagRunId, deadlineTimeGt,
deadlineTimeGte, deadlineTimeLt, deadlineTimeLte, lastUpdatedAtGt,
lastUpdatedAtGte, lastUpdatedAtLt, lastUpdatedAtLte, limit, missed, offset,
orderBy, teams }, queryKey), queryFn: () => DeadlinesService.getDeadlines({
dagId, dagRunId, deadlineTimeGt, deadlineTimeGte, deadline [...]
/**
* Get Dag Deadline Alerts
* Get all deadline alerts defined on a Dag.
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
index d9aa328bcde..6fd74cd8483 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
@@ -1912,10 +1912,11 @@ export const useDashboardServiceDagStatsSuspense =
<TData = Common.DashboardServ
* @param data.lastUpdatedAtGt
* @param data.lastUpdatedAtLte
* @param data.lastUpdatedAtLt
+* @param data.teams
* @returns DeadlineCollectionResponse Successful Response
* @throws ApiError
*/
-export const useDeadlinesServiceGetDeadlinesSuspense = <TData =
Common.DeadlinesServiceGetDeadlinesDefaultResponse, TError = unknown, TQueryKey
extends Array<unknown> = unknown[]>({ dagId, dagRunId, deadlineTimeGt,
deadlineTimeGte, deadlineTimeLt, deadlineTimeLte, lastUpdatedAtGt,
lastUpdatedAtGte, lastUpdatedAtLt, lastUpdatedAtLte, limit, missed, offset,
orderBy }: {
+export const useDeadlinesServiceGetDeadlinesSuspense = <TData =
Common.DeadlinesServiceGetDeadlinesDefaultResponse, TError = unknown, TQueryKey
extends Array<unknown> = unknown[]>({ dagId, dagRunId, deadlineTimeGt,
deadlineTimeGte, deadlineTimeLt, deadlineTimeLte, lastUpdatedAtGt,
lastUpdatedAtGte, lastUpdatedAtLt, lastUpdatedAtLte, limit, missed, offset,
orderBy, teams }: {
dagId: string;
dagRunId: string;
deadlineTimeGt?: string;
@@ -1930,7 +1931,8 @@ export const useDeadlinesServiceGetDeadlinesSuspense =
<TData = Common.Deadlines
missed?: boolean;
offset?: number;
orderBy?: string[];
-}, queryKey?: TQueryKey, options?: Omit<UseQueryOptions<TData, TError>,
"queryKey" | "queryFn">) => useSuspenseQuery<TData, TError>({ queryKey:
Common.UseDeadlinesServiceGetDeadlinesKeyFn({ dagId, dagRunId, deadlineTimeGt,
deadlineTimeGte, deadlineTimeLt, deadlineTimeLte, lastUpdatedAtGt,
lastUpdatedAtGte, lastUpdatedAtLt, lastUpdatedAtLte, limit, missed, offset,
orderBy }, queryKey), queryFn: () => DeadlinesService.getDeadlines({ dagId,
dagRunId, deadlineTimeGt, deadlineTimeGte, deadlin [...]
+ teams?: string[];
+}, queryKey?: TQueryKey, options?: Omit<UseQueryOptions<TData, TError>,
"queryKey" | "queryFn">) => useSuspenseQuery<TData, TError>({ queryKey:
Common.UseDeadlinesServiceGetDeadlinesKeyFn({ dagId, dagRunId, deadlineTimeGt,
deadlineTimeGte, deadlineTimeLt, deadlineTimeLte, lastUpdatedAtGt,
lastUpdatedAtGte, lastUpdatedAtLt, lastUpdatedAtLte, limit, missed, offset,
orderBy, teams }, queryKey), queryFn: () => DeadlinesService.getDeadlines({
dagId, dagRunId, deadlineTimeGt, deadlineTimeGte, [...]
/**
* Get Dag Deadline Alerts
* Get all deadline alerts defined on a Dag.
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 09e4248333c..3198d4162e7 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
@@ -10205,6 +10205,17 @@ export const $DeadlineResponse = {
}
],
title: 'Alert Name'
+ },
+ team_name: {
+ anyOf: [
+ {
+ type: 'string'
+ },
+ {
+ type: 'null'
+ }
+ ],
+ title: 'Team Name'
}
},
type: 'object',
diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
index e90b107ca6f..f5a53a28d22 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
@@ -4893,6 +4893,7 @@ export class DeadlinesService {
* @param data.lastUpdatedAtGt
* @param data.lastUpdatedAtLte
* @param data.lastUpdatedAtLt
+ * @param data.teams
* @returns DeadlineCollectionResponse Successful Response
* @throws ApiError
*/
@@ -4916,7 +4917,8 @@ export class DeadlinesService {
last_updated_at_gte: data.lastUpdatedAtGte,
last_updated_at_gt: data.lastUpdatedAtGt,
last_updated_at_lte: data.lastUpdatedAtLte,
- last_updated_at_lt: data.lastUpdatedAtLt
+ last_updated_at_lt: data.lastUpdatedAtLt,
+ teams: data.teams
},
errors: {
400: 'Bad Request',
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 16810c2219d..ae6b2c2074b 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
@@ -2541,6 +2541,7 @@ export type DeadlineResponse = {
dag_run_id: string;
alert_id?: string | null;
alert_name?: string | null;
+ team_name?: string | null;
};
/**
@@ -4740,6 +4741,7 @@ export type GetDeadlinesData = {
* Attributes to order by, multi criteria sort is supported. Prefix with
`-` for descending order. Supported attributes: `id, deadline_time, created_at,
last_updated_at, missed, dag_id, dag_run_id, alert_name`
*/
orderBy?: Array<(string)>;
+ teams?: Array<(string)>;
};
export type GetDeadlinesResponse = DeadlineCollectionResponse;
diff --git a/airflow-core/src/airflow/ui/src/pages/Deadlines/index.tsx
b/airflow-core/src/airflow/ui/src/pages/Deadlines/index.tsx
index d01d7b33ff6..8f4c7946e1d 100644
--- a/airflow-core/src/airflow/ui/src/pages/Deadlines/index.tsx
+++ b/airflow-core/src/airflow/ui/src/pages/Deadlines/index.tsx
@@ -28,14 +28,16 @@ import { DataTable } from "src/components/DataTable";
import { useTableURLState } from "src/components/DataTable/useTableUrlState";
import { ErrorAlert } from "src/components/ErrorAlert";
import { FilterBar } from "src/components/FilterBar";
+import { TeamName } from "src/components/TeamName";
import Time from "src/components/Time";
import { TruncatedText } from "src/components/TruncatedText";
import { SearchParamsKeys } from "src/constants/searchParams";
+import { useConfig } from "src/queries/useConfig";
import { useDocumentTitle, useFiltersHandler, type FilterableSearchParamsKeys
} from "src/utils";
type DeadlineRow = { row: { original: DeadlineResponse } };
-const createColumns = (translate: TFunction):
Array<ColumnDef<DeadlineResponse>> => [
+const createColumns = (translate: TFunction, multiTeam: boolean):
Array<ColumnDef<DeadlineResponse>> => [
{
accessorKey: "dag_id",
cell: ({ row: { original } }: DeadlineRow) => (
@@ -47,6 +49,16 @@ const createColumns = (translate: TFunction):
Array<ColumnDef<DeadlineResponse>>
),
header: translate("common:dagId"),
},
+ ...(multiTeam
+ ? [
+ {
+ accessorKey: "team_name",
+ cell: ({ row: { original } }: DeadlineRow) => <TeamName
teamName={original.team_name} />,
+ enableSorting: false,
+ header: translate("common:dagDetails.team"),
+ },
+ ]
+ : []),
{
accessorKey: "dag_run_id",
cell: ({ row: { original } }: DeadlineRow) => (
@@ -100,15 +112,18 @@ const deadlinesFilterKeys:
Array<FilterableSearchParamsKeys> = [
export const Deadlines = () => {
const { t: translate } = useTranslation(["browse", "common"]);
+ const multiTeamEnabled = Boolean(useConfig("multi_team"));
useDocumentTitle(translate("common:browse.deadlines"));
const { setTableURLState, tableURLState } = useTableURLState();
const [searchParams] = useSearchParams();
- const { filterConfigs, handleFiltersChange, initialValues } =
useFiltersHandler(deadlinesFilterKeys);
+ const { filterConfigs, handleFiltersChange, initialValues } =
useFiltersHandler(
+ multiTeamEnabled ? [...deadlinesFilterKeys, SearchParamsKeys.TEAMS] :
deadlinesFilterKeys,
+ );
- const columns = createColumns(translate);
+ const columns = createColumns(translate, multiTeamEnabled);
const { pagination, sorting } = tableURLState;
const [sort] = sorting;
@@ -118,6 +133,7 @@ export const Deadlines = () => {
const filteredMissed = searchParams.get(SearchParamsKeys.MISSED);
const deadlineTimeGte = searchParams.get(SearchParamsKeys.DEADLINE_TIME_GTE);
const deadlineTimeLte = searchParams.get(SearchParamsKeys.DEADLINE_TIME_LTE);
+ const teams = searchParams.getAll(SearchParamsKeys.TEAMS);
const missedFilter = filteredMissed === "true" ? true : filteredMissed ===
"false" ? false : undefined;
@@ -130,6 +146,7 @@ export const Deadlines = () => {
missed: missedFilter,
offset: pagination.pageIndex * pagination.pageSize,
orderBy,
+ teams: teams.length > 0 ? teams : undefined,
});
return (
diff --git
a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py
b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py
index c2a8528afa6..8edfcb27816 100644
--- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py
+++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py
@@ -25,7 +25,7 @@ from unittest import mock
import pytest
import time_machine
from fastapi.testclient import TestClient
-from sqlalchemy import delete, func, select, update
+from sqlalchemy import func, select, update
from airflow import plugins_manager
from airflow._shared.module_loading import qualname
@@ -64,6 +64,7 @@ from tests_common.test_utils.db import (
)
from tests_common.test_utils.format_datetime import from_datetime_to_zulu,
from_datetime_to_zulu_without_ms
from tests_common.test_utils.taskinstance import run_task_instance
+from tests_common.test_utils.team import attach_dag_to_team
from unit.listeners.class_listener import ClassBasedListener
if TYPE_CHECKING:
@@ -334,35 +335,6 @@ def get_dag_run_dict(run: DagRun):
}
-def _attach_dag_to_team(session, dag_id: str, *, bundle_name: str, team_name:
str) -> str:
- """
- Associate a Dag with a team via a team-scoped bundle for multi-team tests.
-
- Returns the Dag's original bundle name so the caller can restore it during
cleanup
- (``DagModel.bundle_name`` is a foreign key with no ``ON DELETE`` action).
- """
- original_bundle_name =
session.scalar(select(DagModel.bundle_name).where(DagModel.dag_id == dag_id))
- bundle = DagBundleModel(name=bundle_name)
- bundle.teams.append(Team(name=team_name))
- session.add(bundle)
- session.flush()
- session.execute(update(DagModel).where(DagModel.dag_id ==
dag_id).values(bundle_name=bundle_name))
- session.commit()
- return original_bundle_name
-
-
-def _detach_dag_from_team(
- session, dag_id: str, *, bundle_name: str, team_name: str,
original_bundle_name: str
-) -> None:
- """Undo :func:`_attach_dag_to_team`, restoring the Dag's original
bundle."""
- session.execute(
- update(DagModel).where(DagModel.dag_id ==
dag_id).values(bundle_name=original_bundle_name)
- )
- session.execute(delete(DagBundleModel).where(DagBundleModel.name ==
bundle_name))
- session.execute(delete(Team).where(Team.name == team_name))
- session.commit()
-
-
class TestGetDagRun:
@pytest.mark.parametrize(
("dag_id", "run_id", "state", "run_type", "triggered_by",
"dag_run_note"),
@@ -416,21 +388,10 @@ class TestGetDagRun:
@conf_vars({("core", "multi_team"): "True"})
@pytest.mark.usefixtures("configure_git_connection_for_dag_bundle")
def test_get_dag_run_includes_team_name(self, test_client, session):
- original_bundle_name = _attach_dag_to_team(
- session, DAG1_ID, bundle_name="team-bundle-run",
team_name="team-run"
- )
- try:
+ with attach_dag_to_team(session, DAG1_ID,
bundle_name="team-bundle-run", team_name="team-run"):
response =
test_client.get(f"/dags/{DAG1_ID}/dagRuns/{DAG1_RUN1_ID}")
assert response.status_code == 200
assert response.json()["team_name"] == "team-run"
- finally:
- _detach_dag_from_team(
- session,
- DAG1_ID,
- bundle_name="team-bundle-run",
- team_name="team-run",
- original_bundle_name=original_bundle_name,
- )
def test_get_dag_run_not_found(self, test_client):
response = test_client.get(f"/dags/{DAG1_ID}/dagRuns/invalid")
@@ -489,31 +450,17 @@ class TestGetDagRuns:
@conf_vars({("core", "multi_team"): "True"})
@pytest.mark.usefixtures("configure_git_connection_for_dag_bundle")
def test_get_dag_runs_includes_team_name(self, test_client, session):
- original_bundle_name = _attach_dag_to_team(
- session, DAG1_ID, bundle_name="team-bundle-runs",
team_name="team-runs"
- )
- try:
+ with attach_dag_to_team(session, DAG1_ID,
bundle_name="team-bundle-runs", team_name="team-runs"):
response = test_client.get(f"/dags/{DAG1_ID}/dagRuns")
assert response.status_code == 200
body = response.json()
assert body["dag_runs"]
assert all(run["team_name"] == "team-runs" for run in
body["dag_runs"])
- finally:
- _detach_dag_from_team(
- session,
- DAG1_ID,
- bundle_name="team-bundle-runs",
- team_name="team-runs",
- original_bundle_name=original_bundle_name,
- )
@conf_vars({("core", "multi_team"): "True"})
@pytest.mark.usefixtures("configure_git_connection_for_dag_bundle")
def test_get_dag_runs_filtered_by_team(self, test_client, session):
- original_bundle_name = _attach_dag_to_team(
- session, DAG1_ID, bundle_name="team-bundle-filter",
team_name="team-filter"
- )
- try:
+ with attach_dag_to_team(session, DAG1_ID,
bundle_name="team-bundle-filter", team_name="team-filter"):
response = test_client.get("/dags/~/dagRuns", params={"teams":
["team-filter"]})
assert response.status_code == 200
body = response.json()
@@ -524,14 +471,6 @@ class TestGetDagRuns:
response = test_client.get("/dags/~/dagRuns", params={"teams":
["nonexistent-team"]})
assert response.status_code == 200
assert response.json()["total_entries"] == 0
- finally:
- _detach_dag_from_team(
- session,
- DAG1_ID,
- bundle_name="team-bundle-filter",
- team_name="team-filter",
- original_bundle_name=original_bundle_name,
- )
def test_invalid_order_by_raises_400(self, test_client):
response = test_client.get("/dags/test_dag1/dagRuns?order_by=invalid")
diff --git
a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
index 3784a57f062..4e7e4142015 100644
---
a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
+++
b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
@@ -65,6 +65,7 @@ from tests_common.test_utils.db import (
from tests_common.test_utils.logs import check_last_log
from tests_common.test_utils.mock_operators import MockOperator
from tests_common.test_utils.taskinstance import create_task_instance
+from tests_common.test_utils.team import attach_dag_to_team
pytestmark = pytest.mark.db_test
@@ -76,35 +77,6 @@ DEFAULT_DATETIME_1 =
dt.datetime.fromisoformat(DEFAULT_DATETIME_STR_1)
DEFAULT_DATETIME_2 = dt.datetime.fromisoformat(DEFAULT_DATETIME_STR_2)
-def _attach_dag_to_team(session, dag_id: str, *, bundle_name: str, team_name:
str) -> str:
- """
- Associate a Dag with a team via a team-scoped bundle for multi-team tests.
-
- Returns the Dag's original bundle name so the caller can restore it during
cleanup
- (``DagModel.bundle_name`` is a foreign key with no ``ON DELETE`` action).
- """
- original_bundle_name =
session.scalar(select(DagModel.bundle_name).where(DagModel.dag_id == dag_id))
- bundle = DagBundleModel(name=bundle_name)
- bundle.teams.append(Team(name=team_name))
- session.add(bundle)
- session.flush()
- session.execute(update(DagModel).where(DagModel.dag_id ==
dag_id).values(bundle_name=bundle_name))
- session.commit()
- return original_bundle_name
-
-
-def _detach_dag_from_team(
- session, dag_id: str, *, bundle_name: str, team_name: str,
original_bundle_name: str
-) -> None:
- """Undo :func:`_attach_dag_to_team`, restoring the Dag's original
bundle."""
- session.execute(
- update(DagModel).where(DagModel.dag_id ==
dag_id).values(bundle_name=original_bundle_name)
- )
- session.execute(delete(DagBundleModel).where(DagBundleModel.name ==
bundle_name))
- session.execute(delete(Team).where(Team.name == team_name))
- session.commit()
-
-
class TestTaskInstanceEndpoint:
@staticmethod
def clear_db():
@@ -277,23 +249,14 @@ class TestGetTaskInstance(TestTaskInstanceEndpoint):
@conf_vars({("core", "multi_team"): "True"})
def test_should_include_team_name(self, test_client, session):
self.create_task_instances(session)
- original_bundle_name = _attach_dag_to_team(
+ with attach_dag_to_team(
session, "example_python_operator", bundle_name="team-bundle-ti",
team_name="team-ti"
- )
- try:
+ ):
response = test_client.get(
"/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context"
)
assert response.status_code == 200
assert response.json()["team_name"] == "team-ti"
- finally:
- _detach_dag_from_team(
- session,
- "example_python_operator",
- bundle_name="team-bundle-ti",
- team_name="team-ti",
- original_bundle_name=original_bundle_name,
- )
def test_should_respond_200_with_decorator(self, test_client, session):
self.create_task_instances(session, "example_python_decorator")
@@ -727,26 +690,17 @@ class TestGetMappedTaskInstance(TestTaskInstanceEndpoint):
@conf_vars({("core", "multi_team"): "True"})
def test_should_include_team_name(self, test_client, session):
self.create_task_instances(session)
- original_bundle_name = _attach_dag_to_team(
+ with attach_dag_to_team(
session,
"example_python_operator",
bundle_name="team-bundle-mapped-ti",
team_name="team-mapped-ti",
- )
- try:
+ ):
response = test_client.get(
"/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context/-1",
)
assert response.status_code == 200
assert response.json()["team_name"] == "team-mapped-ti"
- finally:
- _detach_dag_from_team(
- session,
- "example_python_operator",
- bundle_name="team-bundle-mapped-ti",
- team_name="team-mapped-ti",
- original_bundle_name=original_bundle_name,
- )
class TestGetMappedTaskInstances:
@@ -2344,34 +2298,24 @@ class TestGetTaskInstances(TestTaskInstanceEndpoint):
@conf_vars({("core", "multi_team"): "True"})
def test_should_include_team_name(self, test_client, session):
self.create_task_instances(session)
- original_bundle_name = _attach_dag_to_team(
+ with attach_dag_to_team(
session, "example_python_operator", bundle_name="team-bundle-tis",
team_name="team-tis"
- )
- try:
+ ):
response =
test_client.get(f"/dags/{'example_python_operator'}/dagRuns/~/taskInstances")
assert response.status_code == 200
body = response.json()
assert body["task_instances"]
assert all(ti["team_name"] == "team-tis" for ti in
body["task_instances"])
- finally:
- _detach_dag_from_team(
- session,
- "example_python_operator",
- bundle_name="team-bundle-tis",
- team_name="team-tis",
- original_bundle_name=original_bundle_name,
- )
@conf_vars({("core", "multi_team"): "True"})
def test_should_filter_by_team(self, test_client, session):
self.create_task_instances(session)
- original_bundle_name = _attach_dag_to_team(
+ with attach_dag_to_team(
session,
"example_python_operator",
bundle_name="team-bundle-tis-filter",
team_name="team-tis-filter",
- )
- try:
+ ):
response = test_client.get(
"/dags/~/dagRuns/~/taskInstances", params={"teams":
["team-tis-filter"]}
)
@@ -2386,14 +2330,6 @@ class TestGetTaskInstances(TestTaskInstanceEndpoint):
)
assert response.status_code == 200
assert response.json()["total_entries"] == 0
- finally:
- _detach_dag_from_team(
- session,
- "example_python_operator",
- bundle_name="team-bundle-tis-filter",
- team_name="team-tis-filter",
- original_bundle_name=original_bundle_name,
- )
class TestGetTaskDependencies(TestTaskInstanceEndpoint):
diff --git
a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_deadlines.py
b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_deadlines.py
index 369729118c6..dd10f3ded99 100644
--- a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_deadlines.py
+++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_deadlines.py
@@ -37,13 +37,17 @@ from airflow.utils.state import DagRunState
from airflow.utils.types import DagRunTriggeredByType, DagRunType
from tests_common.test_utils.asserts import assert_queries_count
+from tests_common.test_utils.config import conf_vars
from tests_common.test_utils.db import (
+ clear_db_dag_bundles,
clear_db_dags,
clear_db_deadline,
clear_db_deadline_alert,
clear_db_runs,
clear_db_serialized_dags,
+ clear_db_teams,
)
+from tests_common.test_utils.team import attach_dag_to_team
pytestmark = pytest.mark.db_test
@@ -72,13 +76,21 @@ def _cb() -> AsyncCallback:
return AsyncCallback(_CALLBACK_PATH)
[email protected](autouse=True)
-def setup(dag_maker, session):
+def _clear_db():
clear_db_deadline()
clear_db_deadline_alert()
clear_db_runs()
clear_db_dags()
clear_db_serialized_dags()
+ # attach_dag_to_team commits its bundle and team before yielding, so a run
interrupted
+ # inside that context manager leaves rows behind that would collide on the
next run.
+ clear_db_dag_bundles()
+ clear_db_teams()
+
+
[email protected](autouse=True)
+def setup(dag_maker, session):
+ _clear_db()
with dag_maker(DAG_ID, serialized=True, session=session):
EmptyOperator(task_id="task")
@@ -226,11 +238,7 @@ def setup(dag_maker, session):
dag_maker.sync_dagbag_to_db()
session.commit()
yield
- clear_db_deadline()
- clear_db_deadline_alert()
- clear_db_runs()
- clear_db_dags()
- clear_db_serialized_dags()
+ _clear_db()
class TestGetDagRunDeadlines:
@@ -456,6 +464,47 @@ class TestGetDeadlines:
unlinked = [dl for dl in deadlines if dl["alert_name"] is None]
assert all(dl["alert_id"] is None for dl in unlinked)
+ @pytest.mark.parametrize(
+ ("multi_team", "expected_team_name", "expected_queries"),
+ [
+ pytest.param("True", "team-deadlines", 4, id="multi_team_enabled"),
+ pytest.param("False", None, 3, id="multi_team_disabled"),
+ ],
+ )
+ def test_team_name_only_exposed_in_multi_team_mode(
+ self, test_client, session, multi_team, expected_team_name,
expected_queries
+ ):
+ with (
+ conf_vars({("core", "multi_team"): multi_team}),
+ attach_dag_to_team(
+ session, DAG_ID, bundle_name="team-bundle-deadlines",
team_name="team-deadlines"
+ ),
+ ):
+ with assert_queries_count(expected_queries):
+ response = test_client.get("/dags/~/dagRuns/~/deadlines")
+ assert response.status_code == 200
+ teams_by_dag = {dl["dag_id"]: dl["team_name"] for dl in
response.json()["deadlines"]}
+ assert teams_by_dag == {DAG_ID: expected_team_name, DAG_ID_2: None}
+
+ @conf_vars({("core", "multi_team"): "True"})
+ def test_filter_by_team(self, test_client, session):
+ with attach_dag_to_team(
+ session, DAG_ID, bundle_name="team-bundle-deadline-filter",
team_name="team-deadline-filter"
+ ):
+ with assert_queries_count(4):
+ response = test_client.get(
+ "/dags/~/dagRuns/~/deadlines", params={"teams":
["team-deadline-filter"]}
+ )
+ assert response.status_code == 200
+ data = response.json()
+ assert data["total_entries"] == 7
+ assert {dl["dag_id"] for dl in data["deadlines"]} == {DAG_ID}
+
+ # A team with no Dags returns nothing.
+ response = test_client.get("/dags/~/dagRuns/~/deadlines",
params={"teams": ["unknown-team"]})
+ assert response.status_code == 200
+ assert response.json()["total_entries"] == 0
+
def test_filter_nonexistent_dag_returns_empty(self, test_client):
"""Filtering by a dag_id that doesn't exist returns an empty list."""
response = test_client.get("/dags/nonexistent_dag/dagRuns/~/deadlines")
diff --git a/devel-common/src/tests_common/test_utils/team.py
b/devel-common/src/tests_common/test_utils/team.py
new file mode 100644
index 00000000000..9328bca698b
--- /dev/null
+++ b/devel-common/src/tests_common/test_utils/team.py
@@ -0,0 +1,68 @@
+# 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 contextlib import contextmanager
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select, update
+
+if TYPE_CHECKING:
+ from collections.abc import Iterator
+
+ from sqlalchemy.orm import Session
+
+
+@contextmanager
+def attach_dag_to_team(session: Session, dag_id: str, *, bundle_name: str,
team_name: str) -> Iterator[None]:
+ """
+ Associate a Dag with a team through a team-scoped bundle, for multi-team
tests.
+
+ On exit the Dag is moved back to the bundle it started in before the
bundle and team
+ created here are dropped, because ``DagModel.bundle_name`` is a foreign
key with no
+ ``ON DELETE`` action.
+
+ :param session: session used for both the setup and the teardown writes
+ :param dag_id: Dag to move under the team-scoped bundle
+ :param bundle_name: name of the bundle to create; must not already exist
+ :param team_name: name of the team to create and attach the bundle to
+ """
+ from airflow.models.dag import DagModel, clear_team_name_cache
+ from airflow.models.dagbundle import DagBundleModel
+ from airflow.models.team import Team
+
+ original_bundle_name =
session.scalar(select(DagModel.bundle_name).where(DagModel.dag_id == dag_id))
+ bundle = DagBundleModel(name=bundle_name)
+ bundle.teams.append(Team(name=team_name))
+ session.add(bundle)
+ session.flush()
+ session.execute(update(DagModel).where(DagModel.dag_id ==
dag_id).values(bundle_name=bundle_name))
+ session.commit()
+ # DagModel.get_team_name caches by dag_id, so a lookup made before this
re-association
+ # would otherwise keep resolving to the Dag's previous team.
+ clear_team_name_cache()
+ try:
+ yield
+ finally:
+ session.execute(
+ update(DagModel).where(DagModel.dag_id ==
dag_id).values(bundle_name=original_bundle_name)
+ )
+ session.execute(delete(DagBundleModel).where(DagBundleModel.name ==
bundle_name))
+ session.execute(delete(Team).where(Team.name == team_name))
+ session.commit()
+ clear_team_name_cache()
