This is an automated email from the ASF dual-hosted git repository.
pierrejeambrun 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 fe8d8702f22 Paginate partitioned Dag run modal requests (#68412)
fe8d8702f22 is described below
commit fe8d8702f223a33f39a9a49d8520047fbead32d6
Author: Steve Ahn <[email protected]>
AuthorDate: Mon Jun 29 08:29:06 2026 -0700
Paginate partitioned Dag run modal requests (#68412)
* Paginate partitioned Dag run modal requests
* Add stable secondary sort to partitioned Dag run pagination
* Fix asset_expressions typing in partitioned Dag run list response
Build the response via model_validate (matching the rest of the route) so
the
strictly-typed asset_expressions field passes mypy after the main merge; the
direct constructor tripped the static check while keeping identical
behavior.
---
.../api_fastapi/core_api/openapi/_private_ui.yaml | 16 ++
.../core_api/routes/ui/partitioned_dag_runs.py | 21 ++-
.../src/airflow/ui/openapi-gen/queries/common.ts | 6 +-
.../ui/openapi-gen/queries/ensureQueryData.ts | 8 +-
.../src/airflow/ui/openapi-gen/queries/prefetch.ts | 8 +-
.../src/airflow/ui/openapi-gen/queries/queries.ts | 8 +-
.../src/airflow/ui/openapi-gen/queries/suspense.ts | 8 +-
.../ui/openapi-gen/requests/services.gen.ts | 4 +
.../airflow/ui/openapi-gen/requests/types.gen.ts | 2 +
.../pages/DagsList/PartitionScheduleModal.test.tsx | 199 +++++++++++++++++++++
.../src/pages/DagsList/PartitionScheduleModal.tsx | 29 ++-
.../routes/ui/test_partitioned_dag_runs.py | 63 ++++++-
12 files changed, 346 insertions(+), 26 deletions(-)
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 5a62197dcfc..618b158c0f4 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
@@ -152,6 +152,22 @@ paths:
- OAuth2PasswordBearer: []
- HTTPBearer: []
parameters:
+ - name: limit
+ in: query
+ required: false
+ schema:
+ type: integer
+ minimum: 0
+ default: 50
+ title: Limit
+ - name: offset
+ in: query
+ required: false
+ schema:
+ type: integer
+ minimum: 0
+ default: 0
+ title: Offset
- name: dag_id
in: query
required: false
diff --git
a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/partitioned_dag_runs.py
b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/partitioned_dag_runs.py
index 29b393070e3..ad319ced8ee 100644
---
a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/partitioned_dag_runs.py
+++
b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/partitioned_dag_runs.py
@@ -22,8 +22,10 @@ import structlog
from fastapi import Depends, HTTPException, status
from sqlalchemy import and_, select
-from airflow.api_fastapi.common.db.common import SessionDep,
apply_filters_to_select
+from airflow.api_fastapi.common.db.common import SessionDep,
apply_filters_to_select, paginated_select
from airflow.api_fastapi.common.parameters import (
+ QueryLimit,
+ QueryOffset,
QueryPartitionedDagRunDagIdFilter,
QueryPartitionedDagRunHasCreatedDagRunIdFilter,
)
@@ -244,6 +246,8 @@ def _build_response(row, required_count: int,
received_count: int) -> Partitione
)
def get_partitioned_dag_runs(
session: SessionDep,
+ limit: QueryLimit,
+ offset: QueryOffset,
readable_dags_filter: ReadableDagsFilterDep,
dag_id: QueryPartitionedDagRunDagIdFilter,
has_created_dag_run_id: QueryPartitionedDagRunHasCreatedDagRunIdFilter,
@@ -268,14 +272,21 @@ def get_partitioned_dag_runs(
readable_dag_ids = readable_dags_filter.value
if readable_dag_ids is not None:
query =
query.where(AssetPartitionDagRun.target_dag_id.in_(readable_dag_ids))
- query = query.order_by(AssetPartitionDagRun.created_at.desc())
+ query = query.order_by(AssetPartitionDagRun.created_at.desc(),
AssetPartitionDagRun.id.desc())
+
+ query, total_entries = paginated_select(
+ statement=query,
+ offset=offset,
+ limit=limit,
+ session=session,
+ )
if not (rows := session.execute(query).all()):
- if dag_id.value is not None:
+ if dag_id.value is not None and total_entries == 0:
dag_exists =
session.scalar(select(DagModel.dag_id).where(DagModel.dag_id == dag_id.value))
if dag_exists is None:
raise HTTPException(status.HTTP_404_NOT_FOUND, f"Dag with id
{dag_id.value} was not found")
- return PartitionedDagRunCollectionResponse(partitioned_dag_runs=[],
total=0)
+ return PartitionedDagRunCollectionResponse(partitioned_dag_runs=[],
total=total_entries)
# Batch-fetch DagModels (for cached partition_mapper_info), required
assets,
# and APDR log entries in three single queries instead of N per-Dag
queries.
@@ -336,7 +347,7 @@ def get_partitioned_dag_runs(
model_data: dict[str, Any] = {
"partitioned_dag_runs": results,
- "total": len(results),
+ "total": total_entries,
"asset_expressions": asset_expressions,
}
return PartitionedDagRunCollectionResponse.model_validate(model_data)
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 b72cc90eb30..69372a087a3 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts
@@ -891,10 +891,12 @@ export const UseAuthLinksServiceGetCurrentUserInfoKeyFn =
(queryKey?: Array<unkn
export type PartitionedDagRunServiceGetPartitionedDagRunsDefaultResponse =
Awaited<ReturnType<typeof PartitionedDagRunService.getPartitionedDagRuns>>;
export type PartitionedDagRunServiceGetPartitionedDagRunsQueryResult<TData =
PartitionedDagRunServiceGetPartitionedDagRunsDefaultResponse, TError = unknown>
= UseQueryResult<TData, TError>;
export const usePartitionedDagRunServiceGetPartitionedDagRunsKey =
"PartitionedDagRunServiceGetPartitionedDagRuns";
-export const UsePartitionedDagRunServiceGetPartitionedDagRunsKeyFn = ({ dagId,
hasCreatedDagRunId }: {
+export const UsePartitionedDagRunServiceGetPartitionedDagRunsKeyFn = ({ dagId,
hasCreatedDagRunId, limit, offset }: {
dagId?: string;
hasCreatedDagRunId?: boolean;
-} = {}, queryKey?: Array<unknown>) =>
[usePartitionedDagRunServiceGetPartitionedDagRunsKey, ...(queryKey ?? [{ dagId,
hasCreatedDagRunId }])];
+ limit?: number;
+ offset?: number;
+} = {}, queryKey?: Array<unknown>) =>
[usePartitionedDagRunServiceGetPartitionedDagRunsKey, ...(queryKey ?? [{ dagId,
hasCreatedDagRunId, limit, offset }])];
export type PartitionedDagRunServiceGetPendingPartitionedDagRunDefaultResponse
= Awaited<ReturnType<typeof
PartitionedDagRunService.getPendingPartitionedDagRun>>;
export type
PartitionedDagRunServiceGetPendingPartitionedDagRunQueryResult<TData =
PartitionedDagRunServiceGetPendingPartitionedDagRunDefaultResponse, TError =
unknown> = UseQueryResult<TData, TError>;
export const usePartitionedDagRunServiceGetPendingPartitionedDagRunKey =
"PartitionedDagRunServiceGetPendingPartitionedDagRun";
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 74a833b2ccf..ff0174b77e9 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
@@ -1794,15 +1794,19 @@ export const
ensureUseAuthLinksServiceGetCurrentUserInfoData = (queryClient: Que
* Get Partitioned Dag Runs
* Return PartitionedDagRuns. Filter by dag_id and/or has_created_dag_run_id.
* @param data The data for the request.
+* @param data.limit
+* @param data.offset
* @param data.dagId
* @param data.hasCreatedDagRunId
* @returns PartitionedDagRunCollectionResponse Successful Response
* @throws ApiError
*/
-export const ensureUsePartitionedDagRunServiceGetPartitionedDagRunsData =
(queryClient: QueryClient, { dagId, hasCreatedDagRunId }: {
+export const ensureUsePartitionedDagRunServiceGetPartitionedDagRunsData =
(queryClient: QueryClient, { dagId, hasCreatedDagRunId, limit, offset }: {
dagId?: string;
hasCreatedDagRunId?: boolean;
-} = {}) => queryClient.ensureQueryData({ queryKey:
Common.UsePartitionedDagRunServiceGetPartitionedDagRunsKeyFn({ dagId,
hasCreatedDagRunId }), queryFn: () =>
PartitionedDagRunService.getPartitionedDagRuns({ dagId, hasCreatedDagRunId })
});
+ limit?: number;
+ offset?: number;
+} = {}) => queryClient.ensureQueryData({ queryKey:
Common.UsePartitionedDagRunServiceGetPartitionedDagRunsKeyFn({ dagId,
hasCreatedDagRunId, limit, offset }), queryFn: () =>
PartitionedDagRunService.getPartitionedDagRuns({ dagId, hasCreatedDagRunId,
limit, offset }) });
/**
* Get Pending Partitioned Dag Run
* Return full details for pending PartitionedDagRun.
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 9db58a9676c..1f8db278d9d 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
@@ -1794,15 +1794,19 @@ export const
prefetchUseAuthLinksServiceGetCurrentUserInfo = (queryClient: Query
* Get Partitioned Dag Runs
* Return PartitionedDagRuns. Filter by dag_id and/or has_created_dag_run_id.
* @param data The data for the request.
+* @param data.limit
+* @param data.offset
* @param data.dagId
* @param data.hasCreatedDagRunId
* @returns PartitionedDagRunCollectionResponse Successful Response
* @throws ApiError
*/
-export const prefetchUsePartitionedDagRunServiceGetPartitionedDagRuns =
(queryClient: QueryClient, { dagId, hasCreatedDagRunId }: {
+export const prefetchUsePartitionedDagRunServiceGetPartitionedDagRuns =
(queryClient: QueryClient, { dagId, hasCreatedDagRunId, limit, offset }: {
dagId?: string;
hasCreatedDagRunId?: boolean;
-} = {}) => queryClient.prefetchQuery({ queryKey:
Common.UsePartitionedDagRunServiceGetPartitionedDagRunsKeyFn({ dagId,
hasCreatedDagRunId }), queryFn: () =>
PartitionedDagRunService.getPartitionedDagRuns({ dagId, hasCreatedDagRunId })
});
+ limit?: number;
+ offset?: number;
+} = {}) => queryClient.prefetchQuery({ queryKey:
Common.UsePartitionedDagRunServiceGetPartitionedDagRunsKeyFn({ dagId,
hasCreatedDagRunId, limit, offset }), queryFn: () =>
PartitionedDagRunService.getPartitionedDagRuns({ dagId, hasCreatedDagRunId,
limit, offset }) });
/**
* Get Pending Partitioned Dag Run
* Return full details for pending PartitionedDagRun.
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 c10c9635706..c5ec74825be 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
@@ -1794,15 +1794,19 @@ export const useAuthLinksServiceGetCurrentUserInfo =
<TData = Common.AuthLinksSe
* Get Partitioned Dag Runs
* Return PartitionedDagRuns. Filter by dag_id and/or has_created_dag_run_id.
* @param data The data for the request.
+* @param data.limit
+* @param data.offset
* @param data.dagId
* @param data.hasCreatedDagRunId
* @returns PartitionedDagRunCollectionResponse Successful Response
* @throws ApiError
*/
-export const usePartitionedDagRunServiceGetPartitionedDagRuns = <TData =
Common.PartitionedDagRunServiceGetPartitionedDagRunsDefaultResponse, TError =
unknown, TQueryKey extends Array<unknown> = unknown[]>({ dagId,
hasCreatedDagRunId }: {
+export const usePartitionedDagRunServiceGetPartitionedDagRuns = <TData =
Common.PartitionedDagRunServiceGetPartitionedDagRunsDefaultResponse, TError =
unknown, TQueryKey extends Array<unknown> = unknown[]>({ dagId,
hasCreatedDagRunId, limit, offset }: {
dagId?: string;
hasCreatedDagRunId?: boolean;
-} = {}, queryKey?: TQueryKey, options?: Omit<UseQueryOptions<TData, TError>,
"queryKey" | "queryFn">) => useQuery<TData, TError>({ queryKey:
Common.UsePartitionedDagRunServiceGetPartitionedDagRunsKeyFn({ dagId,
hasCreatedDagRunId }, queryKey), queryFn: () =>
PartitionedDagRunService.getPartitionedDagRuns({ dagId, hasCreatedDagRunId })
as TData, ...options });
+ limit?: number;
+ offset?: number;
+} = {}, queryKey?: TQueryKey, options?: Omit<UseQueryOptions<TData, TError>,
"queryKey" | "queryFn">) => useQuery<TData, TError>({ queryKey:
Common.UsePartitionedDagRunServiceGetPartitionedDagRunsKeyFn({ dagId,
hasCreatedDagRunId, limit, offset }, queryKey), queryFn: () =>
PartitionedDagRunService.getPartitionedDagRuns({ dagId, hasCreatedDagRunId,
limit, offset }) as TData, ...options });
/**
* Get Pending Partitioned Dag Run
* Return full details for pending PartitionedDagRun.
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 8c4f6707176..90f75e1cfb3 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
@@ -1794,15 +1794,19 @@ export const
useAuthLinksServiceGetCurrentUserInfoSuspense = <TData = Common.Aut
* Get Partitioned Dag Runs
* Return PartitionedDagRuns. Filter by dag_id and/or has_created_dag_run_id.
* @param data The data for the request.
+* @param data.limit
+* @param data.offset
* @param data.dagId
* @param data.hasCreatedDagRunId
* @returns PartitionedDagRunCollectionResponse Successful Response
* @throws ApiError
*/
-export const usePartitionedDagRunServiceGetPartitionedDagRunsSuspense = <TData
= Common.PartitionedDagRunServiceGetPartitionedDagRunsDefaultResponse, TError =
unknown, TQueryKey extends Array<unknown> = unknown[]>({ dagId,
hasCreatedDagRunId }: {
+export const usePartitionedDagRunServiceGetPartitionedDagRunsSuspense = <TData
= Common.PartitionedDagRunServiceGetPartitionedDagRunsDefaultResponse, TError =
unknown, TQueryKey extends Array<unknown> = unknown[]>({ dagId,
hasCreatedDagRunId, limit, offset }: {
dagId?: string;
hasCreatedDagRunId?: boolean;
-} = {}, queryKey?: TQueryKey, options?: Omit<UseQueryOptions<TData, TError>,
"queryKey" | "queryFn">) => useSuspenseQuery<TData, TError>({ queryKey:
Common.UsePartitionedDagRunServiceGetPartitionedDagRunsKeyFn({ dagId,
hasCreatedDagRunId }, queryKey), queryFn: () =>
PartitionedDagRunService.getPartitionedDagRuns({ dagId, hasCreatedDagRunId })
as TData, ...options });
+ limit?: number;
+ offset?: number;
+} = {}, queryKey?: TQueryKey, options?: Omit<UseQueryOptions<TData, TError>,
"queryKey" | "queryFn">) => useSuspenseQuery<TData, TError>({ queryKey:
Common.UsePartitionedDagRunServiceGetPartitionedDagRunsKeyFn({ dagId,
hasCreatedDagRunId, limit, offset }, queryKey), queryFn: () =>
PartitionedDagRunService.getPartitionedDagRuns({ dagId, hasCreatedDagRunId,
limit, offset }) as TData, ...options });
/**
* Get Pending Partitioned Dag Run
* Return full details for pending PartitionedDagRun.
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 e51b6cf770c..ebb30ffa68c 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
@@ -4663,6 +4663,8 @@ export class PartitionedDagRunService {
* Get Partitioned Dag Runs
* Return PartitionedDagRuns. Filter by dag_id and/or
has_created_dag_run_id.
* @param data The data for the request.
+ * @param data.limit
+ * @param data.offset
* @param data.dagId
* @param data.hasCreatedDagRunId
* @returns PartitionedDagRunCollectionResponse Successful Response
@@ -4673,6 +4675,8 @@ export class PartitionedDagRunService {
method: 'GET',
url: '/ui/partitioned_dag_runs',
query: {
+ limit: data.limit,
+ offset: data.offset,
dag_id: data.dagId,
has_created_dag_run_id: data.hasCreatedDagRunId
},
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 83e95991808..c070c9d74d7 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
@@ -4509,6 +4509,8 @@ export type GenerateTokenResponse2 =
GenerateTokenResponse;
export type GetPartitionedDagRunsData = {
dagId?: string | null;
hasCreatedDagRunId?: boolean | null;
+ limit?: number;
+ offset?: number;
};
export type GetPartitionedDagRunsResponse =
PartitionedDagRunCollectionResponse;
diff --git
a/airflow-core/src/airflow/ui/src/pages/DagsList/PartitionScheduleModal.test.tsx
b/airflow-core/src/airflow/ui/src/pages/DagsList/PartitionScheduleModal.test.tsx
new file mode 100644
index 00000000000..929d0917ab6
--- /dev/null
+++
b/airflow-core/src/airflow/ui/src/pages/DagsList/PartitionScheduleModal.test.tsx
@@ -0,0 +1,199 @@
+/*!
+ * 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 "@testing-library/jest-dom/vitest";
+import { fireEvent, render, screen, waitFor } from "@testing-library/react";
+import type { ReactNode } from "react";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+import type * as OpenapiQueries from "openapi/queries";
+import type { PartitionedDagRunResponse } from "openapi/requests/types.gen";
+import type { TableState } from "src/components/DataTable/types";
+import type * as Ui from "src/components/ui";
+import { Wrapper } from "src/utils/Wrapper";
+
+import { PartitionScheduleModal } from "./PartitionScheduleModal";
+
+vi.mock("react-i18next", () => ({
+ useTranslation: () => ({
+ // eslint-disable-next-line id-length
+ t: (key: string, options?: { count?: number }) =>
+ options?.count === undefined ? key : `${key}:${options.count}`,
+ }),
+}));
+
+vi.mock("src/queries/useConfig", () => ({
+ useConfig: () => 2,
+}));
+
+vi.mock("src/components/ui", async (importOriginal) => {
+ const actual = await importOriginal<typeof Ui>();
+ const DialogPart = ({ children }: { readonly children?: ReactNode }) =>
<div>{children}</div>;
+
+ return {
+ ...actual,
+ Dialog: {
+ ...actual.Dialog,
+ Body: DialogPart,
+ CloseTrigger: () => undefined,
+ Content: DialogPart,
+ Header: DialogPart,
+ Root: ({
+ children,
+ onOpenChange,
+ open,
+ }: {
+ readonly children?: ReactNode;
+ readonly onOpenChange?: () => void;
+ readonly open?: boolean;
+ }) =>
+ open ? (
+ <div>
+ <button onClick={onOpenChange} type="button">
+ close dialog
+ </button>
+ {children}
+ </div>
+ ) : undefined,
+ },
+ };
+});
+
+vi.mock("src/components/DataTable", () => ({
+ DataTable: ({
+ data,
+ initialState,
+ onStateChange,
+ total,
+ }: {
+ readonly data: ReadonlyArray<PartitionedDagRunResponse>;
+ readonly initialState?: TableState;
+ readonly onStateChange?: (state: TableState) => void;
+ readonly total?: number;
+ }) => (
+ <div>
+ <div data-testid="partitioned-dag-run-count">{data.length}</div>
+ <div data-testid="partitioned-dag-run-total">{total}</div>
+ <button
+ onClick={() => {
+ if (initialState !== undefined) {
+ onStateChange?.({
+ ...initialState,
+ pagination: {
+ ...initialState.pagination,
+ pageIndex: 1,
+ },
+ });
+ }
+ }}
+ type="button"
+ >
+ next page
+ </button>
+ </div>
+ ),
+}));
+
+vi.mock("openapi/queries", async (importOriginal) => {
+ const actual = await importOriginal<typeof OpenapiQueries>();
+
+ return {
+ ...actual,
+ usePartitionedDagRunServiceGetPartitionedDagRuns: vi.fn(),
+ };
+});
+
+const { usePartitionedDagRunServiceGetPartitionedDagRuns } = await
import("openapi/queries");
+
+const partitionedDagRun = {
+ dag_id: "example_dag",
+ id: 1,
+ partition_key: "2024-06-01",
+ total_received: 0,
+ total_required: 1,
+} satisfies PartitionedDagRunResponse;
+
+const getPartitionedDagRunsResponse = (partitionedDagRuns:
Array<PartitionedDagRunResponse>) =>
+ ({
+ data: {
+ asset_expressions: null,
+ partitioned_dag_runs: partitionedDagRuns,
+ total: 3,
+ },
+ error: null,
+ isFetching: false,
+ isLoading: false,
+ }) as ReturnType<typeof usePartitionedDagRunServiceGetPartitionedDagRuns>;
+
+describe("PartitionScheduleModal", () => {
+ beforeEach(() => {
+ vi.clearAllMocks();
+
vi.mocked(usePartitionedDagRunServiceGetPartitionedDagRuns).mockReturnValue(
+ getPartitionedDagRunsResponse([partitionedDagRun]),
+ );
+ });
+
+ it("requests partitioned Dag runs for the current table page", async () => {
+ const onClose = vi.fn();
+
+ render(<PartitionScheduleModal dagId="example_dag" onClose={onClose} open
/>, { wrapper: Wrapper });
+
+
expect(usePartitionedDagRunServiceGetPartitionedDagRuns).toHaveBeenLastCalledWith(
+ {
+ dagId: "example_dag",
+ hasCreatedDagRunId: false,
+ limit: 2,
+ offset: 0,
+ },
+ undefined,
+ { enabled: true },
+ );
+
expect(screen.getByTestId("partitioned-dag-run-total")).toHaveTextContent("3");
+
+ fireEvent.click(screen.getByRole("button", { name: "next page" }));
+
+ await waitFor(() =>
+
expect(usePartitionedDagRunServiceGetPartitionedDagRuns).toHaveBeenLastCalledWith(
+ {
+ dagId: "example_dag",
+ hasCreatedDagRunId: false,
+ limit: 2,
+ offset: 2,
+ },
+ undefined,
+ { enabled: true },
+ ),
+ );
+
+ fireEvent.click(screen.getByRole("button", { name: "close dialog" }));
+
+ expect(onClose).toHaveBeenCalledTimes(1);
+ await waitFor(() =>
+
expect(usePartitionedDagRunServiceGetPartitionedDagRuns).toHaveBeenLastCalledWith(
+ {
+ dagId: "example_dag",
+ hasCreatedDagRunId: false,
+ limit: 2,
+ offset: 0,
+ },
+ undefined,
+ { enabled: true },
+ ),
+ );
+ });
+});
diff --git
a/airflow-core/src/airflow/ui/src/pages/DagsList/PartitionScheduleModal.tsx
b/airflow-core/src/airflow/ui/src/pages/DagsList/PartitionScheduleModal.tsx
index ba0ef1a0a13..8a4901ceecb 100644
--- a/airflow-core/src/airflow/ui/src/pages/DagsList/PartitionScheduleModal.tsx
+++ b/airflow-core/src/airflow/ui/src/pages/DagsList/PartitionScheduleModal.tsx
@@ -18,6 +18,7 @@
*/
import { Heading, HStack, Text } from "@chakra-ui/react";
import type { ColumnDef } from "@tanstack/react-table";
+import { useState } from "react";
import { useTranslation } from "react-i18next";
import { FiDatabase } from "react-icons/fi";
@@ -25,9 +26,11 @@ import { usePartitionedDagRunServiceGetPartitionedDagRuns }
from "openapi/querie
import type { PartitionedDagRunResponse } from "openapi/requests/types.gen";
import { AssetProgressCell } from "src/components/AssetProgressCell";
import { DataTable } from "src/components/DataTable";
+import type { TableState } from "src/components/DataTable/types";
import { ErrorAlert } from "src/components/ErrorAlert";
import Time from "src/components/Time";
import { Dialog } from "src/components/ui";
+import { useConfig } from "src/queries/useConfig";
type PartitionScheduleModalProps = {
readonly dagId: string;
@@ -71,9 +74,24 @@ const getColumns = (
export const PartitionScheduleModal = ({ dagId, onClose, open }:
PartitionScheduleModalProps) => {
const { t: translate } = useTranslation("common");
+ const pageSize = (useConfig("fallback_page_limit") as number | undefined) ??
100;
+ const [pageIndex, setPageIndex] = useState(0);
+ const tableState = {
+ pagination: {
+ pageIndex,
+ pageSize,
+ },
+ sorting: [],
+ } satisfies TableState;
+ const { pagination } = tableState;
const { data, error, isFetching, isLoading } =
usePartitionedDagRunServiceGetPartitionedDagRuns(
- { dagId, hasCreatedDagRunId: false },
+ {
+ dagId,
+ hasCreatedDagRunId: false,
+ limit: pagination.pageSize,
+ offset: pagination.pageIndex * pagination.pageSize,
+ },
undefined,
{ enabled: open },
);
@@ -82,8 +100,13 @@ export const PartitionScheduleModal = ({ dagId, onClose,
open }: PartitionSchedu
const total = data?.total ?? 0;
const columns = getColumns(translate, dagId);
+ const handleOpenChange = () => {
+ setPageIndex(0);
+ onClose();
+ };
+
return (
- <Dialog.Root lazyMount onOpenChange={onClose} open={open}
scrollBehavior="inside" unmountOnExit>
+ <Dialog.Root lazyMount onOpenChange={handleOpenChange} open={open}
scrollBehavior="inside" unmountOnExit>
<Dialog.Content backdrop>
<Dialog.Header>
<HStack>
@@ -97,9 +120,11 @@ export const PartitionScheduleModal = ({ dagId, onClose,
open }: PartitionSchedu
<DataTable
columns={columns}
data={partitionedDagRuns}
+ initialState={tableState}
isFetching={isFetching}
isLoading={isLoading}
modelName="partitionedDagRun"
+ onStateChange={(state) => setPageIndex(state.pagination.pageIndex)}
showRowCountHeading={false}
total={total}
/>
diff --git
a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_partitioned_dag_runs.py
b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_partitioned_dag_runs.py
index 3bea3b87d01..93cab212edc 100644
---
a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_partitioned_dag_runs.py
+++
b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_partitioned_dag_runs.py
@@ -68,7 +68,7 @@ class TestGetPartitionedDagRuns:
dag_maker.create_dagrun()
dag_maker.sync_dagbag_to_db()
- with assert_queries_count(3):
+ with assert_queries_count(4):
resp =
test_client.get("/partitioned_dag_runs?dag_id=normal&has_created_dag_run_id=false")
assert resp.status_code == 200
assert resp.json() == {"partitioned_dag_runs": [], "total": 0,
"asset_expressions": None}
@@ -155,16 +155,17 @@ class TestGetPartitionedDagRuns:
)
session.commit()
- # Five batch queries, each independent of len(rows) / num_assets:
- # 1. outer SELECT of AssetPartitionDagRun rows
- # 2. SELECT DagModel WHERE dag_id IN (unique_dag_ids)
- # 3. _fetch_active_assets_per_dag — joined SELECT on AssetModel +
DagScheduleAssetReference
- # 4. SELECT PartitionedAssetKeyLog WHERE asset_partition_dag_run_id IN
(apdr_ids)
- # 5. SELECT timetable for asset_expressions (cached per response)
+ # Six batch queries, each independent of len(rows) / num_assets:
+ # 1. SELECT COUNT of matching AssetPartitionDagRun rows
+ # 2. outer SELECT of paginated AssetPartitionDagRun rows
+ # 3. SELECT DagModel WHERE dag_id IN (unique_dag_ids)
+ # 4. _fetch_active_assets_per_dag — joined SELECT on AssetModel +
DagScheduleAssetReference
+ # 5. SELECT PartitionedAssetKeyLog WHERE asset_partition_dag_run_id IN
(apdr_ids)
+ # 6. SELECT timetable for asset_expressions (cached per response)
# The count does not scale with num_assets or len(rows), so the bump
- # from main's baseline of 3 (3 → 4 → 5 across this branch) is constant
+ # from main's baseline of 3 (3 → 4 → 5 → 6 across this branch) is
constant
# per request, not per row.
- with assert_queries_count(5):
+ with assert_queries_count(6):
resp = test_client.get(
f"/partitioned_dag_runs?dag_id=list_dag"
f"&has_created_dag_run_id={str(has_created_dag_run_id).lower()}"
@@ -178,6 +179,50 @@ class TestGetPartitionedDagRuns:
assert pdr_resp["total_received"] == received_count
assert pdr_resp["total_required"] == num_assets
+ def test_should_paginate_partitioned_dag_runs(self, test_client,
dag_maker, session):
+ with dag_maker(
+ dag_id="paginated_dag",
+
schedule=PartitionedAssetTimetable(assets=Asset(uri="s3://bucket/page",
name="page")),
+ serialized=True,
+ ):
+ EmptyOperator(task_id="t")
+ dag_maker.sync_dagbag_to_db()
+
+ start = pendulum.datetime(2024, 6, 1, tz="UTC")
+ session.add_all(
+ [
+ AssetPartitionDagRun(
+ target_dag_id="paginated_dag",
+ partition_key=f"key-{idx}",
+ created_at=start.add(minutes=idx),
+ )
+ for idx in range(3)
+ ]
+ )
+ session.commit()
+
+ first_page = test_client.get(
+
"/partitioned_dag_runs?dag_id=paginated_dag&has_created_dag_run_id=false&limit=2"
+ )
+ assert first_page.status_code == 200
+ first_body = first_page.json()
+ assert first_body["total"] == 3
+ assert [row["partition_key"] for row in
first_body["partitioned_dag_runs"]] == ["key-2", "key-1"]
+
+ second_page = test_client.get(
+
"/partitioned_dag_runs?dag_id=paginated_dag&has_created_dag_run_id=false&limit=2&offset=2"
+ )
+ assert second_page.status_code == 200
+ second_body = second_page.json()
+ assert second_body["total"] == 3
+ assert [row["partition_key"] for row in
second_body["partitioned_dag_runs"]] == ["key-0"]
+
+ empty_page = test_client.get(
+
"/partitioned_dag_runs?dag_id=paginated_dag&has_created_dag_run_id=false&limit=2&offset=20"
+ )
+ assert empty_page.status_code == 200
+ assert empty_page.json() == {"partitioned_dag_runs": [], "total": 3,
"asset_expressions": None}
+
@pytest.mark.parametrize(
("num_target_assets", "num_other_assets", "received_count"),
[(1, 1, 1), (1, 2, 0), (2, 1, 1)],
