AIP 84: Migrate POST ASSET EVENT legacy API to fast API.

airflow/api_connexion/endpoints/asset_endpoint.py

@@ -323,6 +323,7 @@ def delete_asset_queued_events(
     )
 
 
+@mark_fastapi_migration_done
 @security.requires_access_asset("POST")
 @provide_session
 @action_logging

airflow/api_fastapi/core_api/datamodels/assets.py

@@ -99,3 +99,15 @@ class AssetEventCollectionResponse(BaseModel):
 
     asset_events: list[AssetEventResponse]
     total_entries: int
+
+
+class CreateAssetEventsBody(BaseModel):
+    """Create asset events request."""
+
+    asset_uri: str
+    extra: dict
+
+    class Config:
+        """Pydantic config."""
+
+        extra = "forbid"

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

@@ -346,6 +346,49 @@ paths:
             application/json:
               schema:
                 $ref: '#/components/schemas/HTTPValidationError'
+    post:
+      tags:
+      - Asset
+      summary: Create Asset Events
+      description: Create asset events.
+      operationId: create_asset_events
+      requestBody:
+        required: true
+        content:
+          application/json:
+            schema:
+              $ref: '#/components/schemas/CreateAssetEventsBody'
+      responses:
+        '200':
+          description: Successful Response
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/AssetEventResponse'
+        '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'
   /public/assets/{uri}:
     get:
       tags:
@@ -4037,6 +4080,21 @@ components:
       - message
       title: ConnectionTestResponse
       description: Connection Test serializer for responses.
+    CreateAssetEventsBody:
+      properties:
+        asset_uri:
+          type: string
+          title: Asset Uri
+        extra:
+          type: object
+          title: Extra
+      additionalProperties: false
+      type: object
+      required:
+      - asset_uri
+      - extra
+      title: CreateAssetEventsBody
+      description: Create asset events request.
     DAGCollectionResponse:
       properties:
         dags:

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

@@ -42,9 +42,13 @@
     AssetEventCollectionResponse,
     AssetEventResponse,
     AssetResponse,
+    CreateAssetEventsBody,
 )
 from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
+from airflow.assets import Asset
+from airflow.assets.manager import asset_manager
 from airflow.models.asset import AssetEvent, AssetModel
+from airflow.utils import timezone
 
 assets_router = AirflowRouter(tags=["Asset"], prefix="/assets")
 
@@ -86,7 +90,7 @@ def get_assets(
 
 @assets_router.get(
     "/events",
-    responses=create_openapi_http_exception_doc([404]),
+    responses=create_openapi_http_exception_doc([401, 403, 404]),
 )
 def get_asset_events(
     limit: QueryLimit,
@@ -134,6 +138,36 @@ def get_asset_events(
     )
 
 
+@assets_router.post(
+    "/events",
+    responses=create_openapi_http_exception_doc([401, 403, 404]),
+)
+def create_asset_events(
+    create_asset_request: CreateAssetEventsBody,
+    session: Annotated[Session, Depends(get_session)],
+) -> AssetEventResponse:
+    """Create asset events."""
+    asset = session.scalar(
+        select(AssetModel).where(AssetModel.uri == create_asset_request.asset_uri).limit(1)
+    )
+    if not asset:
+        raise HTTPException(404, f"Asset with uri: `{create_asset_request.asset_uri}` was not found")
+    timestamp = timezone.utcnow()
+
+    create_asset_request.extra["from_rest_api"] = True
+
+    assets_event = asset_manager.register_asset_change(
+        asset=Asset(uri=create_asset_request.asset_uri),
+        timestamp=timestamp,
+        extra=create_asset_request.extra,
+        session=session,
+    )
+
+    if not assets_event:
+        raise HTTPException(404, f"Asset with uri: `{create_asset_request.asset_uri}` was not found")
+    return AssetEventResponse.model_validate(assets_event, from_attributes=True)
+
+
 @assets_router.get(
     "/{uri:path}",
     responses=create_openapi_http_exception_doc([401, 403, 404]),

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

@@ -1045,6 +1045,9 @@ export const UseXcomServiceGetXcomEntryKeyFn = (
     { dagId, dagRunId, deserialize, mapIndex, stringify, taskId, xcomKey },
   ]),
 ];
+export type AssetServiceCreateAssetEventsMutationResult = Awaited<
+  ReturnType<typeof AssetService.createAssetEvents>
+>;
 export type BackfillServiceCreateBackfillMutationResult = Awaited<
   ReturnType<typeof BackfillService.createBackfill>
 >;

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

@@ -32,6 +32,7 @@ import {
 import {
   BackfillPostBody,
   ConnectionBody,
+  CreateAssetEventsBody,
   DAGPatchBody,
   DAGRunPatchBody,
   DagRunState,
@@ -1694,6 +1695,45 @@ export const useXcomServiceGetXcomEntry = <
       }) as TData,
     ...options,
   });
+/**
+ * Create Asset Events
+ * Create asset events.
+ * @param data The data for the request.
+ * @param data.requestBody
+ * @returns AssetEventResponse Successful Response
+ * @throws ApiError
+ */
+export const useAssetServiceCreateAssetEvents = <
+  TData = Common.AssetServiceCreateAssetEventsMutationResult,
+  TError = unknown,
+  TContext = unknown,
+>(
+  options?: Omit<
+    UseMutationOptions<
+      TData,
+      TError,
+      {
+        requestBody: CreateAssetEventsBody;
+      },
+      TContext
+    >,
+    "mutationFn"
+  >,
+) =>
+  useMutation<
+    TData,
+    TError,
+    {
+      requestBody: CreateAssetEventsBody;
+    },
+    TContext
+  >({
+    mutationFn: ({ requestBody }) =>
+      AssetService.createAssetEvents({
+        requestBody,
+      }) as unknown as Promise<TData>,
+    ...options,
+  });
 /**
  * Create Backfill
  * @param data The data for the request.

airflow/ui/openapi-gen/requests/schemas.gen.ts

@@ -695,6 +695,24 @@ export const $ConnectionTestResponse = {
   description: "Connection Test serializer for responses.",
 } as const;
 
+export const $CreateAssetEventsBody = {
+  properties: {
+    asset_uri: {
+      type: "string",
+      title: "Asset Uri",
+    },
+    extra: {
+      type: "object",
+      title: "Extra",
+    },
+  },
+  additionalProperties: false,
+  type: "object",
+  required: ["asset_uri", "extra"],
+  title: "CreateAssetEventsBody",
+  description: "Create asset events request.",
+} as const;
+
 export const $DAGCollectionResponse = {
   properties: {
     dags: {

airflow/ui/openapi-gen/requests/services.gen.ts

@@ -9,6 +9,8 @@ import type {
   GetAssetsResponse,
   GetAssetEventsData,
   GetAssetEventsResponse,
+  CreateAssetEventsData,
+  CreateAssetEventsResponse,
   GetAssetData,
   GetAssetResponse,
   HistoricalMetricsData,
@@ -214,6 +216,31 @@ export class AssetService {
     });
   }
 
+  /**
+   * Create Asset Events
+   * Create asset events.
+   * @param data The data for the request.
+   * @param data.requestBody
+   * @returns AssetEventResponse Successful Response
+   * @throws ApiError
+   */
+  public static createAssetEvents(
+    data: CreateAssetEventsData,
+  ): CancelablePromise<CreateAssetEventsResponse> {
+    return __request(OpenAPI, {
+      method: "POST",
+      url: "/public/assets/events",
+      body: data.requestBody,
+      mediaType: "application/json",
+      errors: {
+        401: "Unauthorized",
+        403: "Forbidden",
+        404: "Not Found",
+        422: "Validation Error",
+      },
+    });
+  }
+
   /**
    * Get Asset
    * Get an asset.

airflow/ui/openapi-gen/requests/types.gen.ts

@@ -174,6 +174,16 @@ export type ConnectionTestResponse = {
   message: string;
 };
 
+/**
+ * Create asset events request.
+ */
+export type CreateAssetEventsBody = {
+  asset_uri: string;
+  extra: {
+    [key: string]: unknown;
+  };
+};
+
 /**
  * DAG Collection serializer for responses.
  */
@@ -984,6 +994,12 @@ export type GetAssetEventsData = {
 
 export type GetAssetEventsResponse = AssetEventCollectionResponse;
 
+export type CreateAssetEventsData = {
+  requestBody: CreateAssetEventsBody;
+};
+
+export type CreateAssetEventsResponse = AssetEventResponse;
+
 export type GetAssetData = {
   uri: string;
 };
@@ -1503,6 +1519,31 @@ export type $OpenApiTs = {
         422: HTTPValidationError;
       };
     };
+    post: {
+      req: CreateAssetEventsData;
+      res: {
+        /**
+         * Successful Response
+         */
+        200: AssetEventResponse;
+        /**
+         * Unauthorized
+         */
+        401: HTTPExceptionResponse;
+        /**
+         * Forbidden
+         */
+        403: HTTPExceptionResponse;
+        /**
+         * Not Found
+         */
+        404: HTTPExceptionResponse;
+        /**
+         * Validation Error
+         */
+        422: HTTPValidationError;
+      };
+    };
   };
   "/public/assets/{uri}": {
     get: {

tests/api_fastapi/core_api/routes/public/test_assets.py

@@ -17,6 +17,7 @@
 from __future__ import annotations
 
 import urllib
+from unittest import mock
 
 import pytest
 
@@ -459,3 +460,32 @@ def test_should_respond_404(self, test_client):
         )
         assert response.status_code == 404
         assert response.json()["detail"] == "The Asset with uri: `s3://bucket/key` was not found"
+
+
+class TestPostAssetEvents(TestAssets):
+    def test_should_respond_200(self, test_client, session):
+        self.create_assets()
+        event_payload = {"asset_uri": "s3://bucket/key/1", "extra": {"foo": "bar"}}
+        response = test_client.post("/public/assets/events", json=event_payload)
+
+        assert response.status_code == 200
+
+        assert response.json() == {
+            "id": mock.ANY,
+            "asset_id": mock.ANY,
+            "uri": "s3://bucket/key/1",
+            "extra": {"foo": "bar", "from_rest_api": True},
+            "source_task_id": mock.ANY,
+            "source_dag_id": mock.ANY,
+            "source_run_id": mock.ANY,
+            "source_map_index": -1,
+            "created_dagruns": [],
+            "timestamp": mock.ANY,
+        }
+
+    def test_invalid_attr_not_allowed(self, test_client, session):
+        self.create_assets()
+        event_invalid_payload = {"asset_uri": "s3://bucket/key/1", "extra": {"foo": "bar"}, "fake": {}}
+        response = test_client.post("/public/assets/events", json=event_invalid_payload)
+
+        assert response.status_code == 422