AIP-84: Migrating delete queued asset events for DAG to fastAPI

airflow/api_connexion/endpoints/asset_endpoint.py

@@ -222,6 +222,7 @@ def delete_dag_asset_queued_event(
     )
 
 
+@mark_fastapi_migration_done
 @security.requires_access_asset("GET")
 @security.requires_access_dag("GET")
 @provide_session
@@ -250,6 +251,7 @@ def get_dag_asset_queued_events(
     )
 
 
+@mark_fastapi_migration_done
 @security.requires_access_asset("DELETE")
 @security.requires_access_dag("GET")
 @action_logging

airflow/api_fastapi/core_api/datamodels/dags.py

@@ -159,3 +159,18 @@ class DAGTagCollectionResponse(BaseModel):
 
     tags: list[str]
     total_entries: int
+
+
+class QueuedEventResponse(BaseModel):
+    """QueuedEvent serializer for responses.."""
+
+    uri: str
+    dag_id: str
+    created_at: datetime
+
+
+class QueuedEventCollectionResponse(BaseModel):
+    """QueuedEventCollection serializer for responses."""
+
+    queued_events: list[QueuedEventResponse]
+    total_entries: int

airflow/api_fastapi/core_api/openapi/v1-generated.yaml

@@ -1864,6 +1864,108 @@ paths:
             application/json:
               schema:
                 $ref: '#/components/schemas/HTTPValidationError'
+  /public/dags/{dag_id}/assets/queuedEvent:
+    get:
+      tags:
+      - DAG
+      summary: Get Dag Asset Queued Events
+      description: Get queued asset events for a DAG.
+      operationId: get_dag_asset_queued_events
+      parameters:
+      - name: dag_id
+        in: path
+        required: true
+        schema:
+          type: string
+          title: Dag Id
+      - name: before
+        in: query
+        required: false
+        schema:
+          type: string
+          title: Before
+      responses:
+        '200':
+          description: Successful Response
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/QueuedEventCollectionResponse'
+        '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'
+    delete:
+      tags:
+      - DAG
+      summary: Delete Dag Asset Queued Events
+      operationId: delete_dag_asset_queued_events
+      parameters:
+      - name: dag_id
+        in: path
+        required: true
+        schema:
+          type: string
+          title: Dag Id
+      - name: before
+        in: query
+        required: false
+        schema:
+          type: string
+          title: Before
+      responses:
+        '204':
+          description: Successful Response
+        '401':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Unauthorized
+        '403':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Forbidden
+        '400':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Bad Request
+        '404':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Not Found
+        '422':
+          description: Validation Error
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPValidationError'
   /public/eventLogs/{event_log_id}:
     get:
       tags:
@@ -5473,6 +5575,41 @@ components:
       - version
       title: ProviderResponse
       description: Provider serializer for responses.
+    QueuedEventCollectionResponse:
+      properties:
+        queued_events:
+          items:
+            $ref: '#/components/schemas/QueuedEventResponse'
+          type: array
+          title: Queued Events
+        total_entries:
+          type: integer
+          title: Total Entries
+      type: object
+      required:
+      - queued_events
+      - total_entries
+      title: QueuedEventCollectionResponse
+      description: QueuedEventCollection serializer for responses.
+    QueuedEventResponse:
+      properties:
+        uri:
+          type: string
+          title: Uri
+        dag_id:
+          type: string
+          title: Dag Id
+        created_at:
+          type: string
+          format: date-time
+          title: Created At
+      type: object
+      required:
+      - uri
+      - dag_id
+      - created_at
+      title: QueuedEventResponse
+      description: QueuedEvent serializer for responses..
     ReprocessBehavior:
       type: string
       enum:

airflow/api_fastapi/core_api/routes/public/dags.py

@@ -20,7 +20,7 @@
 from typing import Annotated
 
 from fastapi import Depends, HTTPException, Query, Request, Response, status
-from sqlalchemy import select, update
+from sqlalchemy import delete, select, update
 from sqlalchemy.orm import Session
 
 from airflow.api.common import delete_dag as delete_dag_module
@@ -50,10 +50,14 @@
     DAGPatchBody,
     DAGResponse,
     DAGTagCollectionResponse,
+    QueuedEventCollectionResponse,
+    QueuedEventResponse,
 )
 from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
 from airflow.exceptions import AirflowException, DagNotFound
 from airflow.models import DAG, DagModel, DagTag
+from airflow.models.asset import AssetDagRunQueue, AssetModel
+from airflow.utils import timezone
 
 dags_router = AirflowRouter(tags=["DAG"], prefix="/dags")
 
@@ -303,3 +307,76 @@ def delete_dag(
             status.HTTP_409_CONFLICT, f"Task instances of dag with id: '{dag_id}' are still running"
         )
     return Response(status_code=status.HTTP_204_NO_CONTENT)
+
+
+@dags_router.get(
+    "/{dag_id}/assets/queuedEvent",
+    responses=create_openapi_http_exception_doc(
+        [
+            status.HTTP_401_UNAUTHORIZED,
+            status.HTTP_403_FORBIDDEN,
+            status.HTTP_404_NOT_FOUND,
+        ]
+    ),
+)
+def get_dag_asset_queued_events(
+    dag_id: str,
+    session: Annotated[Session, Depends(get_session)],
+    before: str = Query(None),
+) -> QueuedEventCollectionResponse:
+    """Get queued asset events for a DAG."""
+    where_clause = [AssetDagRunQueue.target_dag_id == dag_id]
+    if before:
+        before_parsed = timezone.parse(before)
+        where_clause.append(AssetDagRunQueue.created_at < before_parsed)
+    query = (
+        select(AssetDagRunQueue, AssetModel.uri)
+        .join(AssetModel, AssetDagRunQueue.asset_id == AssetModel.id)
+        .where(*where_clause)
+    )
+    result = session.execute(query).all()
+    total_entries = len(result)
+    if not result:
+        raise HTTPException(status.HTTP_404_NOT_FOUND, f"Queue event with dag_id: `{dag_id}` was not found")
+    queued_events = [
+        QueuedEventResponse(created_at=adrq.created_at, dag_id=adrq.target_dag_id, uri=uri)
+        for adrq, uri in result
+    ]
+    return QueuedEventCollectionResponse(
+        queued_events=[
+            QueuedEventResponse.model_validate(queued_event, from_attributes=True)
+            for queued_event in queued_events
+        ],
+        total_entries=total_entries,
+    )
+
+
+@dags_router.delete(
+    "/{dag_id}/assets/queuedEvent",
+    status_code=status.HTTP_204_NO_CONTENT,
+    responses=create_openapi_http_exception_doc(
+        [
+            status.HTTP_400_BAD_REQUEST,
+            status.HTTP_401_UNAUTHORIZED,
+            status.HTTP_403_FORBIDDEN,
+            status.HTTP_404_NOT_FOUND,
+        ]
+    ),
+)
+def delete_dag_asset_queued_events(
+    dag_id: str,
+    session: Annotated[Session, Depends(get_session)],
+    before: str = Query(None),
+):
+    where_clause = [AssetDagRunQueue.target_dag_id == dag_id]
+    if before:
+        before_parsed = timezone.parse(before)
+        where_clause.append(AssetDagRunQueue.created_at < before_parsed)
+
+    delete_statement = delete(AssetDagRunQueue).where(*where_clause)
+    result = session.execute(delete_statement)
+
+    if result.rowcount == 0:
+        raise HTTPException(
+            status.HTTP_404_NOT_FOUND, detail=f"Queue event with dag_id: `{dag_id}` was not found"
+        )

airflow/ui/openapi-gen/queries/common.ts

@@ -478,6 +478,28 @@ export const UseDagServiceGetDagDetailsKeyFn = (
   },
   queryKey?: Array<unknown>,
 ) => [useDagServiceGetDagDetailsKey, ...(queryKey ?? [{ dagId }])];
+export type DagServiceGetDagAssetQueuedEventsDefaultResponse = Awaited<
+  ReturnType<typeof DagService.getDagAssetQueuedEvents>
+>;
+export type DagServiceGetDagAssetQueuedEventsQueryResult<
+  TData = DagServiceGetDagAssetQueuedEventsDefaultResponse,
+  TError = unknown,
+> = UseQueryResult<TData, TError>;
+export const useDagServiceGetDagAssetQueuedEventsKey =
+  "DagServiceGetDagAssetQueuedEvents";
+export const UseDagServiceGetDagAssetQueuedEventsKeyFn = (
+  {
+    before,
+    dagId,
+  }: {
+    before?: string;
+    dagId: string;
+  },
+  queryKey?: Array<unknown>,
+) => [
+  useDagServiceGetDagAssetQueuedEventsKey,
+  ...(queryKey ?? [{ before, dagId }]),
+];
 export type EventLogServiceGetEventLogDefaultResponse = Awaited<
   ReturnType<typeof EventLogService.getEventLog>
 >;
@@ -1096,6 +1118,9 @@ export type DagRunServiceDeleteDagRunMutationResult = Awaited<
 export type DagServiceDeleteDagMutationResult = Awaited<
   ReturnType<typeof DagService.deleteDag>
 >;
+export type DagServiceDeleteDagAssetQueuedEventsMutationResult = Awaited<
+  ReturnType<typeof DagService.deleteDagAssetQueuedEvents>
+>;
 export type PoolServiceDeletePoolMutationResult = Awaited<
   ReturnType<typeof PoolService.deletePool>
 >;

airflow/ui/openapi-gen/queries/prefetch.ts

@@ -622,6 +622,32 @@ export const prefetchUseDagServiceGetDagDetails = (
     queryKey: Common.UseDagServiceGetDagDetailsKeyFn({ dagId }),
     queryFn: () => DagService.getDagDetails({ dagId }),
   });
+/**
+ * Get Dag Asset Queued Events
+ * Get queued asset events for a DAG.
+ * @param data The data for the request.
+ * @param data.dagId
+ * @param data.before
+ * @returns QueuedEventCollectionResponse Successful Response
+ * @throws ApiError
+ */
+export const prefetchUseDagServiceGetDagAssetQueuedEvents = (
+  queryClient: QueryClient,
+  {
+    before,
+    dagId,
+  }: {
+    before?: string;
+    dagId: string;
+  },
+) =>
+  queryClient.prefetchQuery({
+    queryKey: Common.UseDagServiceGetDagAssetQueuedEventsKeyFn({
+      before,
+      dagId,
+    }),
+    queryFn: () => DagService.getDagAssetQueuedEvents({ before, dagId }),
+  });
 /**
  * Get Event Log
  * @param data The data for the request.