From 7650f09d5647be6331a804c513a80df43be8a10e Mon Sep 17 00:00:00 2001
From: Karthikeyan Singaravelan <tir.karthi@gmail.com>
Date: Wed, 4 Sep 2024 22:17:11 +0530
Subject: [PATCH] Implement pagination for the new home page. (#41905)

* Implement pagination for the home page.

* Add page number buttons.

* Change default number of dags to display to 50.

* Fix imports.

* Fix PR comments.
---
 airflow/ui/src/app.tsx      | 23 ++++++++-
 airflow/ui/src/dagsList.tsx | 97 ++++++++++++++++++++++++++++++++++++-
 2 files changed, 117 insertions(+), 3 deletions(-)

diff --git a/airflow/ui/src/app.tsx b/airflow/ui/src/app.tsx
index d5d582c7902e7..a8e2b213f8ae9 100644
--- a/airflow/ui/src/app.tsx
+++ b/airflow/ui/src/app.tsx
@@ -17,21 +17,40 @@
  * under the License.
  */
 
+import { useState } from "react";
 import { Box, Spinner } from "@chakra-ui/react";
+import { PaginationState } from "@tanstack/react-table";
 
 import { useDagServiceGetDags } from "openapi/queries";
 import { DagsList } from "src/dagsList";
 import { Nav } from "src/nav";
 
 export const App = () => {
-  const { data, isLoading } = useDagServiceGetDags();
+  // TODO: Change this to be taken from airflow.cfg
+  const pageSize = 50;
+  const [pagination, setPagination] = useState<PaginationState>({
+    pageIndex: 0,
+    pageSize: pageSize,
+  });
+
+  const { data, isLoading } = useDagServiceGetDags({
+    limit: pagination.pageSize,
+    offset: pagination.pageIndex * pagination.pageSize,
+  });
 
   return (
     <Box maxWidth="100vw">
       <Nav />
       <Box p={3}>
         {isLoading && <Spinner />}
-        {!isLoading && !!data?.dags && <DagsList data={data.dags} />}
+        {!isLoading && !!data?.dags && (
+          <DagsList
+            data={data.dags}
+            total={data.total_entries}
+            pagination={pagination}
+            setPagination={setPagination}
+          />
+        )}
       </Box>
     </Box>
   );
diff --git a/airflow/ui/src/dagsList.tsx b/airflow/ui/src/dagsList.tsx
index 5e0e68fdf9c1d..b5ac0f5751720 100644
--- a/airflow/ui/src/dagsList.tsx
+++ b/airflow/ui/src/dagsList.tsx
@@ -23,9 +23,13 @@ import {
   useReactTable,
   getCoreRowModel,
   getExpandedRowModel,
+  getPaginationRowModel,
   ColumnDef,
   flexRender,
   Row,
+  OnChangeFn,
+  PaginationState,
+  Table as TanStackTable,
 } from "@tanstack/react-table";
 import { MdExpandMore } from "react-icons/md";
 import {
@@ -33,6 +37,7 @@ import {
   Code,
   Table as ChakraTable,
   Thead,
+  Button,
   Td,
   Th,
   Tr,
@@ -84,16 +89,85 @@ const columns: ColumnDef<DAG>[] = [
 
 type TableProps<TData> = {
   data: TData[];
+  total: number | undefined;
   columns: ColumnDef<TData>[];
   renderSubComponent: (props: { row: Row<TData> }) => React.ReactElement;
   getRowCanExpand: (row: Row<TData>) => boolean;
+  pagination: PaginationState;
+  setPagination: OnChangeFn<PaginationState>;
+};
+
+type PaginatorProps<TData> = {
+  table: TanStackTable<TData>;
+};
+
+const TablePaginator = ({ table }: PaginatorProps<DAG>) => {
+  const pageInterval = 3;
+  const currentPageNumber = table.getState().pagination.pageIndex + 1;
+  const startPageNumber = Math.max(1, currentPageNumber - pageInterval);
+  const endPageNumber = Math.min(
+    table.getPageCount(),
+    startPageNumber + pageInterval * 2
+  );
+  const pageNumbers = [];
+
+  for (let index = startPageNumber; index <= endPageNumber; index++) {
+    pageNumbers.push(
+      <Button
+        borderRadius={0}
+        key={index}
+        isDisabled={index === currentPageNumber}
+        onClick={() => table.setPageIndex(index - 1)}
+      >
+        {index}
+      </Button>
+    );
+  }
+
+  return (
+    <Box mt={2} mb={2}>
+      <Button
+        borderRadius={0}
+        onClick={() => table.firstPage()}
+        isDisabled={!table.getCanPreviousPage()}
+      >
+        {"<<"}
+      </Button>
+
+      <Button
+        borderRadius={0}
+        onClick={() => table.previousPage()}
+        isDisabled={!table.getCanPreviousPage()}
+      >
+        {"<"}
+      </Button>
+      {pageNumbers}
+      <Button
+        borderRadius={0}
+        onClick={() => table.nextPage()}
+        isDisabled={!table.getCanNextPage()}
+      >
+        {">"}
+      </Button>
+      <Button
+        borderRadius={0}
+        onClick={() => table.lastPage()}
+        isDisabled={!table.getCanNextPage()}
+      >
+        {">>"}
+      </Button>
+    </Box>
+  );
 };
 
 const Table = ({
   data,
+  total,
   columns,
   renderSubComponent,
   getRowCanExpand,
+  pagination,
+  setPagination,
 }: TableProps<DAG>) => {
   const table = useReactTable<DAG>({
     data,
@@ -101,6 +175,13 @@ const Table = ({
     getRowCanExpand,
     getCoreRowModel: getCoreRowModel(),
     getExpandedRowModel: getExpandedRowModel(),
+    getPaginationRowModel: getPaginationRowModel(),
+    onPaginationChange: setPagination,
+    rowCount: total ?? 0,
+    manualPagination: true,
+    state: {
+      pagination,
+    },
   });
 
   return (
@@ -156,6 +237,7 @@ const Table = ({
           })}
         </Tbody>
       </ChakraTable>
+      <TablePaginator table={table} />
     </TableContainer>
   );
 };
@@ -168,13 +250,26 @@ const renderSubComponent = ({ row }: { row: Row<DAG> }) => {
   );
 };
 
-export const DagsList = ({ data }: { data: DAG[] }) => {
+export const DagsList = ({
+  data,
+  total,
+  pagination,
+  setPagination,
+}: {
+  data: DAG[];
+  total: number | undefined;
+  pagination: PaginationState;
+  setPagination: OnChangeFn<PaginationState>;
+}) => {
   return (
     <Table
       data={data}
+      total={total}
       columns={columns}
       getRowCanExpand={() => true}
       renderSubComponent={renderSubComponent}
+      pagination={pagination}
+      setPagination={setPagination}
     />
   );
 };