AIP-84 Migrate delete Dag Run endpoint to FastAPI (apache#42910)

* delete dag_run init * delete dag_run init * add delete dag run endpoint * add tests for delete dag run
harjeevanmaan · Oct 23, 2024 · da218e5 · da218e5
1 parent 634ae23
commit da218e5
Show file tree

Hide file tree

Showing 8 changed files with 193 additions and 0 deletions.
diff --git a/airflow/api_connexion/endpoints/dag_run_endpoint.py b/airflow/api_connexion/endpoints/dag_run_endpoint.py
@@ -78,6 +78,7 @@
     from airflow.api_connexion.types import APIResponse
 
 
+@mark_fastapi_migration_done
 @security.requires_access_dag("DELETE", DagAccessEntity.RUN)
 @provide_session
 @action_logging

diff --git a/airflow/api_fastapi/openapi/v1-generated.yaml b/airflow/api_fastapi/openapi/v1-generated.yaml
@@ -679,6 +679,58 @@ paths:
             application/json:
               schema:
                 $ref: '#/components/schemas/HTTPValidationError'
+    delete:
+      tags:
+      - DagRun
+      summary: Delete Dag Run
+      description: Delete a DAG Run entry.
+      operationId: delete_dag_run
+      parameters:
+      - name: dag_id
+        in: path
+        required: true
+        schema:
+          type: string
+          title: Dag Id
+      - name: dag_run_id
+        in: path
+        required: true
+        schema:
+          type: string
+          title: Dag Run Id
+      responses:
+        '204':
+          description: Successful Response
+        '400':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Bad Request
+        '401':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Unauthorized
+        '403':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Forbidden
+        '404':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Not Found
+        '422':
+          description: Validation Error
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPValidationError'
 components:
   schemas:
     ConnectionResponse:

diff --git a/airflow/api_fastapi/views/public/dag_run.py b/airflow/api_fastapi/views/public/dag_run.py
@@ -42,3 +42,18 @@ async def get_dag_run(
         )
 
     return DAGRunResponse.model_validate(dag_run, from_attributes=True)
+
+
+@dag_run_router.delete(
+    "/{dag_run_id}", status_code=204, responses=create_openapi_http_exception_doc([400, 401, 403, 404])
+)
+async def delete_dag_run(dag_id: str, dag_run_id: str, session: Annotated[Session, Depends(get_session)]):
+    """Delete a DAG Run entry."""
+    dag_run = session.scalar(select(DagRun).filter_by(dag_id=dag_id, run_id=dag_run_id))
+
+    if dag_run is None:
+        raise HTTPException(
+            404, f"The DagRun with dag_id: `{dag_id}` and run_id: `{dag_run_id}` was not found"
+        )
+
+    session.delete(dag_run)
diff --git a/airflow/ui/openapi-gen/queries/common.ts b/airflow/ui/openapi-gen/queries/common.ts
@@ -197,3 +197,6 @@ export type ConnectionServiceDeleteConnectionMutationResult = Awaited<
 export type VariableServiceDeleteVariableMutationResult = Awaited<
   ReturnType<typeof VariableService.deleteVariable>
 >;
+export type DagRunServiceDeleteDagRunMutationResult = Awaited<
+  ReturnType<typeof DagRunService.deleteDagRun>
+>;
diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts
@@ -506,3 +506,46 @@ export const useVariableServiceDeleteVariable = <
       }) as unknown as Promise<TData>,
     ...options,
   });
+/**
+ * Delete Dag Run
+ * Delete a DAG Run entry.
+ * @param data The data for the request.
+ * @param data.dagId
+ * @param data.dagRunId
+ * @returns void Successful Response
+ * @throws ApiError
+ */
+export const useDagRunServiceDeleteDagRun = <
+  TData = Common.DagRunServiceDeleteDagRunMutationResult,
+  TError = unknown,
+  TContext = unknown,
+>(
+  options?: Omit<
+    UseMutationOptions<
+      TData,
+      TError,
+      {
+        dagId: string;
+        dagRunId: string;
+      },
+      TContext
+    >,
+    "mutationFn"
+  >,
+) =>
+  useMutation<
+    TData,
+    TError,
+    {
+      dagId: string;
+      dagRunId: string;
+    },
+    TContext
+  >({
+    mutationFn: ({ dagId, dagRunId }) =>
+      DagRunService.deleteDagRun({
+        dagId,
+        dagRunId,
+      }) as unknown as Promise<TData>,
+    ...options,
+  });
diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts
@@ -27,6 +27,8 @@ import type {
   GetVariableResponse,
   GetDagRunData,
   GetDagRunResponse,
+  DeleteDagRunData,
+  DeleteDagRunResponse,
 } from "./types.gen";
 
 export class AssetService {
@@ -391,4 +393,33 @@ export class DagRunService {
       },
     });
   }
+
+  /**
+   * Delete Dag Run
+   * Delete a DAG Run entry.
+   * @param data The data for the request.
+   * @param data.dagId
+   * @param data.dagRunId
+   * @returns void Successful Response
+   * @throws ApiError
+   */
+  public static deleteDagRun(
+    data: DeleteDagRunData,
+  ): CancelablePromise<DeleteDagRunResponse> {
+    return __request(OpenAPI, {
+      method: "DELETE",
+      url: "/public/dags/{dag_id}/dagRuns/{dag_run_id}",
+      path: {
+        dag_id: data.dagId,
+        dag_run_id: data.dagRunId,
+      },
+      errors: {
+        400: "Bad Request",
+        401: "Unauthorized",
+        403: "Forbidden",
+        404: "Not Found",
+        422: "Validation Error",
+      },
+    });
+  }
 }
diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts
@@ -355,6 +355,13 @@ export type GetDagRunData = {
 
 export type GetDagRunResponse = DAGRunResponse;
 
+export type DeleteDagRunData = {
+  dagId: string;
+  dagRunId: string;
+};
+
+export type DeleteDagRunResponse = void;
+
 export type $OpenApiTs = {
   "/ui/next_run_assets/{dag_id}": {
     get: {
@@ -657,5 +664,34 @@ export type $OpenApiTs = {
         422: HTTPValidationError;
       };
     };
+    delete: {
+      req: DeleteDagRunData;
+      res: {
+        /**
+         * Successful Response
+         */
+        204: void;
+        /**
+         * Bad Request
+         */
+        400: HTTPExceptionResponse;
+        /**
+         * Unauthorized
+         */
+        401: HTTPExceptionResponse;
+        /**
+         * Forbidden
+         */
+        403: HTTPExceptionResponse;
+        /**
+         * Not Found
+         */
+        404: HTTPExceptionResponse;
+        /**
+         * Validation Error
+         */
+        422: HTTPValidationError;
+      };
+    };
   };
 };
diff --git a/tests/api_fastapi/views/public/test_dag_run.py b/tests/api_fastapi/views/public/test_dag_run.py
@@ -136,3 +136,15 @@ def test_get_dag_run_not_found(test_client):
     assert response.status_code == 404
     body = response.json()
     assert body["detail"] == "The DagRun with dag_id: `test_dag1` and run_id: `invalid` was not found"
+
+
+class TestDeleteDagRun:
+    def test_delete_dag_run(self, test_client):
+        response = test_client.delete(f"/public/dags/{DAG1_ID}/dagRuns/{DAG1_RUN1_ID}")
+        assert response.status_code == 204
+
+    def test_delete_dag_run_not_found(self, test_client):
+        response = test_client.delete(f"/public/dags/{DAG1_ID}/dagRuns/invalid")
+        assert response.status_code == 404
+        body = response.json()
+        assert body["detail"] == "The DagRun with dag_id: `test_dag1` and run_id: `invalid` was not found"