AIP-84: Migrating GET Dataset events for DAG runs api to fastAPI

airflow/api_connexion/endpoints/asset_endpoint.py

@@ -47,6 +47,7 @@
 from airflow.assets.manager import asset_manager
 from airflow.models.asset import AssetDagRunQueue, AssetEvent, AssetModel
 from airflow.utils import timezone
+from airflow.utils.api_migration import mark_fastapi_migration_done
 from airflow.utils.db import get_query_count
 from airflow.utils.session import NEW_SESSION, provide_session
 from airflow.www.decorators import action_logging
@@ -77,6 +78,7 @@ def get_asset(*, uri: str, session: Session = NEW_SESSION) -> APIResponse:
     return asset_schema.dump(asset)
 
 
+@mark_fastapi_migration_done
 @security.requires_access_asset("GET")
 @format_parameters({"limit": check_limit})
 @provide_session

airflow/api_connexion/endpoints/dag_run_endpoint.py

@@ -115,6 +115,7 @@ def get_dag_run(
         raise BadRequest("DAGRunSchema error", detail=str(e))
 
 
+@mark_fastapi_migration_done
 @security.requires_access_dag("GET", DagAccessEntity.RUN)
 @security.requires_access_asset("GET")
 @provide_session

airflow/api_fastapi/common/parameters.py

@@ -29,6 +29,7 @@
 
 from airflow.api_connexion.endpoints.task_instance_endpoint import _convert_ti_states
 from airflow.models import Base, Connection
+from airflow.models.asset import AssetModel, DagScheduleAssetReference, TaskOutletAssetReference
 from airflow.models.dag import DagModel, DagTag
 from airflow.models.dagrun import DagRun
 from airflow.models.dagwarning import DagWarning, DagWarningType
@@ -408,6 +409,37 @@ def depends(self, dag_id: str | None = None) -> _DagIdFilter:
         return self.set_value(dag_id)
 
 
+class _UriPatternSearch(_SearchParam):
+    """Search on uri."""
+
+    def __init__(self, skip_none: bool = True) -> None:
+        super().__init__(AssetModel.uri, skip_none)
+
+    def depends(self, uri_pattern: str | None = None) -> _UriPatternSearch:
+        return self.set_value(uri_pattern)
+
+
+class _DagIdAssetReferenceFilter(BaseParam[list[str]]):
+    """Search on dag_id."""
+
+    def __init__(self, skip_none: bool = True) -> None:
+        super().__init__(AssetModel.consuming_dags, skip_none)
+
+    def depends(self, dag_ids: list[str] = Query(None)) -> _DagIdAssetReferenceFilter:
+        # needed to handle cases where dag_ids=a1,b1
+        if dag_ids and len(dag_ids) == 1 and "," in dag_ids[0]:
+            dag_ids = dag_ids[0].split(",")
+        return self.set_value(dag_ids)
+
+    def to_orm(self, select: Select) -> Select:
+        if self.value is None and self.skip_none:
+            return select
+        return select.where(
+            (AssetModel.consuming_dags.any(DagScheduleAssetReference.dag_id.in_(self.value)))
+            | (AssetModel.producing_tasks.any(TaskOutletAssetReference.dag_id.in_(self.value)))
+        )
+
+
 class Range(BaseModel, Generic[T]):
     """Range with a lower and upper bound."""
 
@@ -493,8 +525,15 @@ def depends_float(
 
 # DAGTags
 QueryDagTagPatternSearch = Annotated[_DagTagNamePatternSearch, Depends(_DagTagNamePatternSearch().depends)]
+
 # TI
 QueryTIStateFilter = Annotated[_TIStateFilter, Depends(_TIStateFilter().depends)]
 QueryTIPoolFilter = Annotated[_TIPoolFilter, Depends(_TIPoolFilter().depends)]
 QueryTIQueueFilter = Annotated[_TIQueueFilter, Depends(_TIQueueFilter().depends)]
 QueryTIExecutorFilter = Annotated[_TIExecutorFilter, Depends(_TIExecutorFilter().depends)]
+
+# Assets
+QueryUriPatternSearch = Annotated[_UriPatternSearch, Depends(_UriPatternSearch().depends)]
+QueryAssetDagIdPatternSearch = Annotated[
+    _DagIdAssetReferenceFilter, Depends(_DagIdAssetReferenceFilter().depends)
+]

airflow/api_fastapi/core_api/datamodels/assets.py

@@ -0,0 +1,103 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from datetime import datetime
+
+from pydantic import BaseModel, Field
+
+
+class DagScheduleAssetReference(BaseModel):
+    """Serializable version of the DagScheduleAssetReference ORM SqlAlchemyModel."""
+
+    dag_id: str
+    created_at: datetime
+    updated_at: datetime
+
+
+class TaskOutletAssetReference(BaseModel):
+    """Serializable version of the TaskOutletAssetReference ORM SqlAlchemyModel."""
+
+    dag_id: str
+    task_id: str
+    created_at: datetime
+    updated_at: datetime
+
+
+class AssetAliasSchema(BaseModel):
+    """Serializable version of the AssetAliasSchema ORM SqlAlchemyModel."""
+
+    id: int
+    name: str
+
+
+class AssetResponse(BaseModel):
+    """Asset serializer for responses."""
+
+    id: int
+    uri: str
+    extra: dict | None = None
+    created_at: datetime
+    updated_at: datetime
+    consuming_dags: list[DagScheduleAssetReference]
+    producing_tasks: list[TaskOutletAssetReference]
+    aliases: list[AssetAliasSchema]
+
+
+class AssetCollectionResponse(BaseModel):
+    """Asset collection response."""
+
+    assets: list[AssetResponse]
+    total_entries: int
+
+
+class DagRunAssetReference(BaseModel):
+    """Serializable version of the DagRunAssetReference ORM SqlAlchemyModel."""
+
+    run_id: str
+    dag_id: str
+    execution_date: datetime = Field(alias="logical_date")
+    start_date: datetime
+    end_date: datetime | None = None
+    state: str
+    data_interval_start: datetime
+    data_interval_end: datetime
+
+
+class AssetEventResponse(BaseModel):
+    """Asset event serializer for responses."""
+
+    id: int
+    asset_id: int
+    # piggyback on the fix for https://github.com/apache/airflow/issues/43845 for asset_uri
+    # meanwhile, unblock by adding uri below
+    uri: str
+    extra: dict | None = None
+    source_task_id: str | None = None
+    source_dag_id: str | None = None
+    source_run_id: str | None = None
+    source_map_index: int
+    created_dagruns: list[DagRunAssetReference]
+    timestamp: datetime
+
+
+class AssetEventCollectionResponse(BaseModel):
+    """Asset collection response."""
+
+    asset_events: list[AssetEventResponse]
+    total_entries: int