Add worker maintenance mode functionality to Edge3 provider UI

providers/edge3/src/airflow/providers/edge3/openapi/v2-edge-generated.yaml

@@ -586,6 +586,67 @@ paths:
       security:
       - OAuth2PasswordBearer: []
       - HTTPBearer: []
+  /edge_worker/ui/worker/{worker_name}/maintenance:
+    post:
+      tags:
+      - UI
+      summary: Request Worker Maintenance
+      description: Put a worker into maintenance mode.
+      operationId: request_worker_maintenance
+      parameters:
+      - name: worker_name
+        in: path
+        required: true
+        schema:
+          type: string
+          title: Worker Name
+      requestBody:
+        required: true
+        content:
+          application/json:
+            schema:
+              $ref: '#/components/schemas/MaintenanceRequest'
+      responses:
+        '200':
+          description: Successful Response
+          content:
+            application/json:
+              schema:
+                type: 'null'
+                title: Response Request Worker Maintenance
+        '422':
+          description: Validation Error
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPValidationError'
+    delete:
+      tags:
+      - UI
+      summary: Exit Worker Maintenance
+      description: Exit a worker from maintenance mode.
+      operationId: exit_worker_maintenance
+      parameters:
+      - name: worker_name
+        in: path
+        required: true
+        schema:
+          type: string
+          title: Worker Name
+      responses:
+        '200':
+          description: Successful Response
+          content:
+            application/json:
+              schema:
+                type: 'null'
+                title: Response Exit Worker Maintenance
+        '422':
+          description: Validation Error
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPValidationError'
 components:
   schemas:
     BundleInfo:
@@ -794,6 +855,17 @@ components:
       - total_entries
       title: JobCollectionResponse
       description: Job Collection serializer.
+    MaintenanceRequest:
+      properties:
+        maintenance_comment:
+          type: string
+          title: Maintenance Comment
+          description: Comment describing the maintenance reason.
+      type: object
+      required:
+      - maintenance_comment
+      title: MaintenanceRequest
+      description: Request body for maintenance operations.
     PushLogsBody:
       properties:
         log_chunk_time:

providers/edge3/src/airflow/providers/edge3/plugins/www/openapi-gen/queries/common.ts

@@ -28,6 +28,8 @@ export const UseUiServiceJobsKeyFn = (queryKey?: Array<unknown>) => [useUiServic
 export type JobsServiceFetchMutationResult = Awaited<ReturnType<typeof JobsService.fetch>>;
 export type LogsServicePushLogsMutationResult = Awaited<ReturnType<typeof LogsService.pushLogs>>;
 export type WorkerServiceRegisterMutationResult = Awaited<ReturnType<typeof WorkerService.register>>;
+export type UiServiceRequestWorkerMaintenanceMutationResult = Awaited<ReturnType<typeof UiService.requestWorkerMaintenance>>;
 export type JobsServiceStateMutationResult = Awaited<ReturnType<typeof JobsService.state>>;
 export type WorkerServiceSetStateMutationResult = Awaited<ReturnType<typeof WorkerService.setState>>;
 export type WorkerServiceUpdateQueuesMutationResult = Awaited<ReturnType<typeof WorkerService.updateQueues>>;
+export type UiServiceExitWorkerMaintenanceMutationResult = Awaited<ReturnType<typeof UiService.exitWorkerMaintenance>>;

providers/edge3/src/airflow/providers/edge3/plugins/www/openapi-gen/queries/queries.ts

@@ -2,7 +2,7 @@
 
 import { UseMutationOptions, UseQueryOptions, useMutation, useQuery } from "@tanstack/react-query";
 import { JobsService, LogsService, MonitorService, UiService, WorkerService } from "../requests/services.gen";
-import { PushLogsBody, TaskInstanceState, WorkerQueueUpdateBody, WorkerQueuesBody, WorkerStateBody } from "../requests/types.gen";
+import { MaintenanceRequest, PushLogsBody, TaskInstanceState, WorkerQueueUpdateBody, WorkerQueuesBody, WorkerStateBody } from "../requests/types.gen";
 import * as Common from "./common";
 export const useLogsServiceLogfilePath = <TData = Common.LogsServiceLogfilePathDefaultResponse, TError = unknown, TQueryKey extends Array<unknown> = unknown[]>({ authorization, dagId, mapIndex, runId, taskId, tryNumber }: {
   authorization: string;
@@ -50,6 +50,13 @@ export const useWorkerServiceRegister = <TData = Common.WorkerServiceRegisterMut
   requestBody: WorkerStateBody;
   workerName: string;
 }, TContext>({ mutationFn: ({ authorization, requestBody, workerName }) => WorkerService.register({ authorization, requestBody, workerName }) as unknown as Promise<TData>, ...options });
+export const useUiServiceRequestWorkerMaintenance = <TData = Common.UiServiceRequestWorkerMaintenanceMutationResult, TError = unknown, TContext = unknown>(options?: Omit<UseMutationOptions<TData, TError, {
+  requestBody: MaintenanceRequest;
+  workerName: string;
+}, TContext>, "mutationFn">) => useMutation<TData, TError, {
+  requestBody: MaintenanceRequest;
+  workerName: string;
+}, TContext>({ mutationFn: ({ requestBody, workerName }) => UiService.requestWorkerMaintenance({ requestBody, workerName }) as unknown as Promise<TData>, ...options });
 export const useJobsServiceState = <TData = Common.JobsServiceStateMutationResult, TError = unknown, TContext = unknown>(options?: Omit<UseMutationOptions<TData, TError, {
   authorization: string;
   dagId: string;
@@ -85,3 +92,8 @@ export const useWorkerServiceUpdateQueues = <TData = Common.WorkerServiceUpdateQ
   requestBody: WorkerQueueUpdateBody;
   workerName: string;
 }, TContext>({ mutationFn: ({ authorization, requestBody, workerName }) => WorkerService.updateQueues({ authorization, requestBody, workerName }) as unknown as Promise<TData>, ...options });
+export const useUiServiceExitWorkerMaintenance = <TData = Common.UiServiceExitWorkerMaintenanceMutationResult, TError = unknown, TContext = unknown>(options?: Omit<UseMutationOptions<TData, TError, {
+  workerName: string;
+}, TContext>, "mutationFn">) => useMutation<TData, TError, {
+  workerName: string;
+}, TContext>({ mutationFn: ({ workerName }) => UiService.exitWorkerMaintenance({ workerName }) as unknown as Promise<TData>, ...options });

providers/edge3/src/airflow/providers/edge3/plugins/www/openapi-gen/requests/schemas.gen.ts

@@ -252,6 +252,20 @@ export const $JobCollectionResponse = {
     description: 'Job Collection serializer.'
 } as const;
 
+export const $MaintenanceRequest = {
+    properties: {
+        maintenance_comment: {
+            type: 'string',
+            title: 'Maintenance Comment',
+            description: 'Comment describing the maintenance reason.'
+        }
+    },
+    type: 'object',
+    required: ['maintenance_comment'],
+    title: 'MaintenanceRequest',
+    description: 'Request body for maintenance operations.'
+} as const;
+
 export const $PushLogsBody = {
     properties: {
         log_chunk_time: {

providers/edge3/src/airflow/providers/edge3/plugins/www/openapi-gen/requests/services.gen.ts

@@ -3,7 +3,7 @@
 import type { CancelablePromise } from './core/CancelablePromise';
 import { OpenAPI } from './core/OpenAPI';
 import { request as __request } from './core/request';
-import type { FetchData, FetchResponse, StateData, StateResponse, LogfilePathData, LogfilePathResponse, PushLogsData, PushLogsResponse, RegisterData, RegisterResponse, SetStateData, SetStateResponse, UpdateQueuesData, UpdateQueuesResponse, HealthResponse, WorkerResponse, JobsResponse } from './types.gen';
+import type { FetchData, FetchResponse, StateData, StateResponse, LogfilePathData, LogfilePathResponse, PushLogsData, PushLogsResponse, RegisterData, RegisterResponse, SetStateData, SetStateResponse, UpdateQueuesData, UpdateQueuesResponse, HealthResponse, WorkerResponse, JobsResponse, RequestWorkerMaintenanceData, RequestWorkerMaintenanceResponse, ExitWorkerMaintenanceData, ExitWorkerMaintenanceResponse } from './types.gen';
 
 export class JobsService {
     /**
@@ -286,4 +286,49 @@ export class UiService {
         });
     }
 
+    /**
+     * Request Worker Maintenance
+     * Put a worker into maintenance mode.
+     * @param data The data for the request.
+     * @param data.workerName
+     * @param data.requestBody
+     * @returns null Successful Response
+     * @throws ApiError
+     */
+    public static requestWorkerMaintenance(data: RequestWorkerMaintenanceData): CancelablePromise<RequestWorkerMaintenanceResponse> {
+        return __request(OpenAPI, {
+            method: 'POST',
+            url: '/edge_worker/ui/worker/{worker_name}/maintenance',
+            path: {
+                worker_name: data.workerName
+            },
+            body: data.requestBody,
+            mediaType: 'application/json',
+            errors: {
+                422: 'Validation Error'
+            }
+        });
+    }
+
+    /**
+     * Exit Worker Maintenance
+     * Exit a worker from maintenance mode.
+     * @param data The data for the request.
+     * @param data.workerName
+     * @returns null Successful Response
+     * @throws ApiError
+     */
+    public static exitWorkerMaintenance(data: ExitWorkerMaintenanceData): CancelablePromise<ExitWorkerMaintenanceResponse> {
+        return __request(OpenAPI, {
+            method: 'DELETE',
+            url: '/edge_worker/ui/worker/{worker_name}/maintenance',
+            path: {
+                worker_name: data.workerName
+            },
+            errors: {
+                422: 'Validation Error'
+            }
+        });
+    }
+
 }

providers/edge3/src/airflow/providers/edge3/plugins/www/openapi-gen/requests/types.gen.ts

@@ -126,6 +126,16 @@ export type JobCollectionResponse = {
     total_entries: number;
 };
 
+/**
+ * Request body for maintenance operations.
+ */
+export type MaintenanceRequest = {
+    /**
+     * Comment describing the maintenance reason.
+     */
+    maintenance_comment: string;
+};
+
 /**
  * Incremental new log content from worker.
  */
@@ -460,6 +470,19 @@ export type WorkerResponse = WorkerCollectionResponse;
 
 export type JobsResponse = JobCollectionResponse;
 
+export type RequestWorkerMaintenanceData = {
+    requestBody: MaintenanceRequest;
+    workerName: string;
+};
+
+export type RequestWorkerMaintenanceResponse = null;
+
+export type ExitWorkerMaintenanceData = {
+    workerName: string;
+};
+
+export type ExitWorkerMaintenanceResponse = null;
+
 export type $OpenApiTs = {
     '/edge_worker/v1/jobs/fetch/{worker_name}': {
         post: {
@@ -652,4 +675,32 @@ export type $OpenApiTs = {
             };
         };
     };
+    '/edge_worker/ui/worker/{worker_name}/maintenance': {
+        post: {
+            req: RequestWorkerMaintenanceData;
+            res: {
+                /**
+                 * Successful Response
+                 */
+                200: null;
+                /**
+                 * Validation Error
+                 */
+                422: HTTPValidationError;
+            };
+        };
+        delete: {
+            req: ExitWorkerMaintenanceData;
+            res: {
+                /**
+                 * Successful Response
+                 */
+                200: null;
+                /**
+                 * Validation Error
+                 */
+                422: HTTPValidationError;
+            };
+        };
+    };
 };