feat(task_sdk): add AssetEventCollectionResponse

Author: Lee-W

airflow/api_fastapi/execution_api/datamodels/asset.py

@@ -76,6 +76,12 @@ def redact_extra(cls, v: dict):
         return redact(v)
 
 
+class AssetEventCollectionResponse(BaseModel):
+    """Collection of AssetEventResponse."""
+
+    asset_events: list[AssetEventResponse]
+
+
 class AssetProfile(BaseModel):
     """
     Profile of an Asset.

airflow/api_fastapi/execution_api/routes/asset_events.py

@@ -19,12 +19,14 @@
 
 from typing import Annotated
 
-from fastapi import HTTPException, Query, status
+from fastapi import Query, status
 from sqlalchemy import and_, select
 
 from airflow.api_fastapi.common.db.common import SessionDep
 from airflow.api_fastapi.common.router import AirflowRouter
-from airflow.api_fastapi.execution_api.datamodels.asset import AssetEventResponse
+from airflow.api_fastapi.execution_api.datamodels.asset import (
+    AssetEventCollectionResponse,
+)
 from airflow.models.asset import AssetAliasModel, AssetEvent, AssetModel
 
 # TODO: Add dependency on JWT token
@@ -36,41 +38,22 @@
 )
 
 
-class LazyAssetEventSelectSequence(LazySelectSequence[AssetEvent]):
-    """
-    List-like interface to lazily access AssetEvent rows.
-
-    :meta private:
-    """
-
-    @staticmethod
-    def _rebuild_select(stmt: TextClause) -> Select:
-        return select(AssetEvent).from_statement(stmt)
-
-    @staticmethod
-    def _process_row(row: Row) -> AssetEvent:
-        return row[0]
-
-
-def _get_asset_event_through_sql_clause(
+def _get_asset_events_through_sql_clauses(
     *, join_clause, where_clause, session: SessionDep
-) -> AssetEventResponse:
-    asset_event = LazyAssetEventSelectSequence.from_select(
-        select(AssetEvent).join(join_clause).where(where_clause),
-        order_by=[AssetEvent.timestamp],
-        session=session,
+) -> AssetEventCollectionResponse:
+    asset_events = session.scalars(
+        select(AssetEvent).join(join_clause).where(where_clause).order_by(AssetEvent.timestamp)
     )
-    _raise_if_not_found(asset_event=asset_event, msg="Not found")
-    return AssetEventResponse.model_validate(asset_event)
+    return AssetEventCollectionResponse.model_validate({"asset_events": asset_events or []})
 
 
 @router.get("/by-asset-name-uri")
 def get_asset_event_by_asset_name_uri(
     name: Annotated[str, Query(description="The name of the Asset")],
     uri: Annotated[str, Query(description="The URI of the Asset")],
     session: SessionDep,
-) -> AssetEventResponse:
-    return _get_asset_event_through_sql_clause(
+) -> AssetEventCollectionResponse:
+    return _get_asset_events_through_sql_clauses(
         join_clause=AssetEvent.asset,
         where_clause=and_(AssetModel.name == name, AssetModel.uri == uri),
         session=session,
@@ -81,8 +64,8 @@ def get_asset_event_by_asset_name_uri(
 def get_asset_event_by_uri(
     uri: Annotated[str, Query(description="The URI of the Asset")],
     session: SessionDep,
-) -> AssetEventResponse:
-    return _get_asset_event_through_sql_clause(
+) -> AssetEventCollectionResponse:
+    return _get_asset_events_through_sql_clauses(
         join_clause=AssetEvent.asset,
         where_clause=and_(AssetModel.uri == uri, AssetModel.active.has()),
         session=session,
@@ -93,8 +76,8 @@ def get_asset_event_by_uri(
 def get_asset_event_by_name(
     name: Annotated[str, Query(description="The name of the Asset")],
     session: SessionDep,
-) -> AssetEventResponse:
-    return _get_asset_event_through_sql_clause(
+) -> AssetEventCollectionResponse:
+    return _get_asset_events_through_sql_clauses(
         join_clause=AssetEvent.asset,
         where_clause=and_(AssetModel.uri == name, AssetModel.active.has()),
         session=session,
@@ -105,20 +88,9 @@ def get_asset_event_by_name(
 def get_asset_event_by_alias_name(
     name: Annotated[str, Query(description="The name of the Asset Alias")],
     session: SessionDep,
-) -> AssetEventResponse:
-    return _get_asset_event_through_sql_clause(
+) -> AssetEventCollectionResponse:
+    return _get_asset_events_through_sql_clauses(
         join_clause=AssetEvent.source_aliases,
         where_clause=(AssetAliasModel.name == name),
         session=session,
     )
-
-
-def _raise_if_not_found(asset_event, msg):
-    if asset_event is None:
-        raise HTTPException(
-            status.HTTP_404_NOT_FOUND,
-            detail={
-                "reason": "not_found",
-                "message": msg,
-            },
-        )

task_sdk/src/airflow/sdk/api/client.py

@@ -35,7 +35,7 @@
 from airflow.api_fastapi.execution_api.datamodels.taskinstance import TIRuntimeCheckPayload
 from airflow.sdk import __version__
 from airflow.sdk.api.datamodels._generated import (
-    AssetEventResponse,
+    AssetEventCollectionResponse,
     AssetResponse,
     ConnectionResponse,
     DagRunType,
@@ -340,7 +340,7 @@ def __init__(self, client: Client):
 
     def get(
         self, name: str | None = None, uri: str | None = None, alias_name: str | None = None
-    ) -> AssetResponse:
+    ) -> AssetEventCollectionResponse:
         """Get Asset value from the API server."""
         if name and uri:
             resp = self.client.get("asset-events/by-asset-name-uri", params={"name": name, "uri": uri})
@@ -353,7 +353,7 @@ def get(
         else:
             raise ValueError("Either `name`, `uri` or `alias_name` must be provided")
 
-        return AssetEventResponse.model_validate_json(resp.read())
+        return AssetEventCollectionResponse.model_validate_json(resp.read())
 
 
 class BearerAuth(httpx.Auth):

task_sdk/src/airflow/sdk/api/datamodels/_generated.py

@@ -1,4 +1,3 @@
-#
 # 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
@@ -370,3 +369,11 @@ class TITerminalStatePayload(BaseModel):
     )
     state: TerminalStateNonSuccess
     end_date: Annotated[datetime, Field(title="End Date")]
+
+
+class AssetEventCollectionResponse(BaseModel):
+    """
+    Collection of AssetEventResponse.
+    """
+
+    asset_events: Annotated[list[AssetEventResponse], Field(title="Asset Events")]

task_sdk/src/airflow/sdk/definitions/asset/__init__.py

@@ -18,6 +18,7 @@
 from __future__ import annotations
 
 import contextlib
+import datetime
 import logging
 import operator
 import os
@@ -694,6 +695,39 @@ def as_expression(self) -> Any:
         return {"all": [o.as_expression() for o in self.objects]}
 
 
+@attrs.define(kw_only=True)
+class DagRunAssetReference:
+    run_id: str
+    dag_id: str
+    start_date: datetime.datetime
+    state: str
+    end_date: datetime.datetime | None = None
+
+    logical_date: datetime.datetime | None = None
+    data_interval_start: datetime.datetime | None = None
+    data_interval_end: datetime.datetime | None = None
+
+
+@attrs.define(kw_only=True)
+class AssetEvent:
+    """Representation of asset event to be triggered by an asset alias."""
+
+    id: int
+    asset_id: int
+
+    created_dagruns: list[DagRunAssetReference]
+    timestamp: datetime.datetime
+
+    uri: str | None = None
+    name: str | None = None
+    group: str | None = None
+    extra: dict[str, Any] | None = None
+    source_task_id: str | None = None
+    source_dag_id: str | None = None
+    source_run_id: str | None = None
+    source_map_index: int | None = None
+
+
 @attrs.define
 class AssetAliasEvent:
     """Representation of asset event to be triggered by an asset alias."""

task_sdk/src/airflow/sdk/execution_time/comms.py

@@ -51,7 +51,7 @@
 from pydantic import BaseModel, ConfigDict, Field, JsonValue
 
 from airflow.sdk.api.datamodels._generated import (
-    AssetEventResponse,
+    AssetEventCollectionResponse,
     AssetResponse,
     BundleInfo,
     ConnectionResponse,
@@ -105,23 +105,28 @@ def from_asset_response(cls, asset_response: AssetResponse) -> AssetResult:
         return cls(**asset_response.model_dump(exclude_defaults=True), type="AssetResult")
 
 
-class AssetEventResult(AssetEventResponse):
-    """Response to ReadXCom request."""
+class AssetEventCollectionResult(AssetEventCollectionResponse):
+    """Response to GetAssetEvent request."""
 
-    type: Literal["AssetEventResult"] = "AssetEventResult"
+    type: Literal["AssetEventCollectionResult"] = "AssetEventCollectionResult"
 
     @classmethod
-    def from_asset_event_response(cls, asset_event_response: AssetEventResponse) -> AssetEventResult:
+    def from_asset_event_collection_response(
+        cls, asset_event_collection_response: AssetEventCollectionResponse
+    ) -> AssetEventCollectionResult:
         """
-        Get AssetEventResult from AssetEventResponse.
+        Get AssetEventCollectionResult from AssetEventCollectionResponse.
 
         AssetEventCollectionResponse is autogenerated from the API schema, so we need to convert it to AssetEventCollectionResult
         for communication between the Supervisor and the task process.
         """
         # Exclude defaults to avoid sending unnecessary data
         # Pass the type as AssetResult explicitly so we can then call model_dump_json with exclude_unset=True
         # to avoid sending unset fields (which are defaults in our case).
-        return cls(**asset_event_response.model_dump(exclude_defaults=True), type="AssetEventResult")
+        return cls(
+            **asset_event_collection_response.model_dump(exclude_defaults=True),
+            type="AssetEventCollectionResult",
+        )
 
 
 class XComResult(XComResponse):
@@ -204,7 +209,7 @@ class OKResponse(BaseModel):
 ToTask = Annotated[
     Union[
         AssetResult,
-        AssetEventResult,
+        AssetEventCollectionResult,
         ConnectionResult,
         ErrorResponse,
         PrevSuccessfulDagRunResult,

task_sdk/src/airflow/sdk/execution_time/context.py

@@ -24,27 +24,21 @@
 import attrs
 import structlog
 
-from airflow.models.asset import AssetEvent
 from airflow.sdk.definitions._internal.contextmanager import _CURRENT_CONTEXT
 from airflow.sdk.definitions._internal.types import NOTSET
 from airflow.sdk.definitions.asset import (
     Asset,
     AssetAlias,
     AssetAliasEvent,
     AssetAliasUniqueKey,
+    AssetEvent,
     AssetNameRef,
     AssetRef,
     AssetUniqueKey,
     AssetUriRef,
     BaseAssetUniqueKey,
 )
 from airflow.sdk.exceptions import AirflowRuntimeError, ErrorType
-from airflow.sdk.execution_time.comms import (
-    AssetEventResult,
-    ErrorResponse,
-    GetAssetEventByAliasName,
-    GetAssetEventByName,
-)
 
 if TYPE_CHECKING:
     from uuid import UUID
@@ -314,16 +308,16 @@ def __init__(self, inlets: list) -> None:
     def __iter__(self) -> Iterator[Asset | AssetAlias]:
         return iter(self._inlets)
 
-            AssetEventCollectionResult,
-            ErrorResponse,
-            GetAssetEventByAliasName,
-            GetAssetEventByName,
     def __len__(self) -> int:
         return len(self._inlets)
 
     def __getitem__(self, key: int | Asset | AssetAlias | AssetRef):
         from airflow.sdk.definitions.asset import Asset
         from airflow.sdk.execution_time.comms import (
+            AssetEventCollectionResult,
+            ErrorResponse,
+            GetAssetEventByAliasName,
+            GetAssetEventByName,
             GetAssetEventByNameUri,
             GetAssetEventByUri,
         )
@@ -362,8 +356,8 @@ def __getitem__(self, key: int | Asset | AssetAlias | AssetRef):
             raise AirflowRuntimeError(msg)
 
         if TYPE_CHECKING:
-            assert isinstance(msg, AssetEventResult)
-        return AssetEvent(**msg.model_dump(excldue={"type"}))
+            assert isinstance(msg, AssetEventCollectionResult)
+        return [AssetEvent(**event) for event in msg.model_dump()["asset_events"]]
 
 
 @cache  # Prevent multiple API access.

task_sdk/src/airflow/sdk/execution_time/supervisor.py

@@ -61,7 +61,7 @@
     VariableResponse,
 )
 from airflow.sdk.execution_time.comms import (
-    AssetEventResult,
+    AssetEventCollectionResult,
     AssetResult,
     ConnectionResult,
     DeferTask,
@@ -831,19 +831,27 @@ def _handle_request(self, msg: ToSupervisor, log: FilteringBoundLogger):
             resp = asset_result.model_dump_json(exclude_unset=True).encode()
         elif isinstance(msg, GetAssetEventByNameUri):
             asset_event_resp = self.client.asset_events.get(uri=msg.uri, name=msg.name)
-            asset_event_result = AssetEventResult.from_asset_event_response(asset_event_resp)
+            asset_event_result = AssetEventCollectionResult.from_asset_event_collection_response(
+                asset_event_resp
+            )
             resp = asset_event_result.model_dump_json(exclude_unset=True).encode()
         elif isinstance(msg, GetAssetEventByName):
             asset_event_resp = self.client.asset_events.get(name=msg.name)
-            asset_event_result = AssetEventResult.from_asset_event_response(asset_event_resp)
+            asset_event_result = AssetEventCollectionResult.from_asset_event_collection_response(
+                asset_event_resp
+            )
             resp = asset_event_result.model_dump_json(exclude_unset=True).encode()
         elif isinstance(msg, GetAssetEventByUri):
             asset_event_resp = self.client.asset_events.get(uri=msg.uri)
-            asset_event_result = AssetEventResult.from_asset_event_response(asset_event_resp)
+            asset_event_result = AssetEventCollectionResult.from_asset_event_collection_response(
+                asset_event_resp
+            )
             resp = asset_event_result.model_dump_json(exclude_unset=True).encode()
         elif isinstance(msg, GetAssetEventByAliasName):
             asset_event_resp = self.client.asset_events.get(name=msg.alias_name)
-            asset_event_result = AssetEventResult.from_asset_event_response(asset_event_resp)
+            asset_event_result = AssetEventCollectionResult.from_asset_event_collection_response(
+                asset_event_resp
+            )
             resp = asset_event_result.model_dump_json(exclude_unset=True).encode()
         elif isinstance(msg, GetPrevSuccessfulDagRun):
             dagrun_resp = self.client.task_instances.get_previous_successful_dagrun(self.id)