fix(task_sdk/context): add _get_asset_events_from_db for fixing tests

Author: Lee-W

airflow/utils/context.py

@@ -29,22 +29,37 @@
 )
 
 import attrs
-from sqlalchemy import select
+from sqlalchemy import and_, select
 
 from airflow.models.asset import (
+    AssetAliasModel,
+    AssetEvent,
     AssetModel,
 )
+from airflow.sdk.definitions.asset import (
+    Asset,
+    AssetAlias,
+    AssetAliasUniqueKey,
+    AssetNameRef,
+    AssetRef,
+    AssetUniqueKey,
+    AssetUriRef,
+)
 from airflow.sdk.definitions.context import Context
 from airflow.sdk.execution_time.context import (
     ConnectionAccessor as ConnectionAccessorSDK,
     InletEventsAccessors as InletEventsAccessorsSDK,
     OutletEventAccessors as OutletEventAccessorsSDK,
     VariableAccessor as VariableAccessorSDK,
 )
+from airflow.utils.db import LazySelectSequence
 from airflow.utils.session import create_session
 from airflow.utils.types import NOTSET
 
 if TYPE_CHECKING:
+    from sqlalchemy.engine import Row
+    from sqlalchemy.sql.expression import Select, TextClause
+
     from airflow.sdk.types import OutletEventAccessorsProtocol
 
 # NOTE: Please keep this in sync with the following:
@@ -150,6 +165,22 @@ def _get_asset_from_db(name: str | None = None, uri: str | None = None) -> Asset
         return asset.to_public()
 
 
+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]
+
+
 @attrs.define(init=False)
 class InletEventsAccessors(InletEventsAccessorsSDK):
     """
@@ -158,6 +189,37 @@ class InletEventsAccessors(InletEventsAccessorsSDK):
     :meta private:
     """
 
+    def _get_asset_events_from_db(self, obj: Asset | AssetAlias | AssetRef):
+        if isinstance(obj, Asset):
+            asset = self._assets[AssetUniqueKey.from_asset(obj)]
+            join_clause = AssetEvent.asset
+            where_clause = and_(AssetModel.name == asset.name, AssetModel.uri == asset.uri)
+        elif isinstance(obj, AssetAlias):
+            asset_alias = self._asset_aliases[AssetAliasUniqueKey.from_asset_alias(obj)]
+            join_clause = AssetEvent.source_aliases
+            where_clause = AssetAliasModel.name == asset_alias.name
+        elif isinstance(obj, AssetNameRef):
+            try:
+                asset = next(a for k, a in self._assets.items() if k.name == obj.name)
+            except StopIteration:
+                raise KeyError(obj) from None
+            join_clause = AssetEvent.asset
+            where_clause = and_(AssetModel.name == asset.name, AssetModel.active.has())
+        elif isinstance(obj, AssetUriRef):
+            try:
+                asset = next(a for k, a in self._assets.items() if k.uri == obj.uri)
+            except StopIteration:
+                raise KeyError(obj) from None
+            join_clause = AssetEvent.asset
+            where_clause = and_(AssetModel.uri == asset.uri, AssetModel.active.has())
+
+        with create_session() as session:
+            return LazyAssetEventSelectSequence.from_select(
+                select(AssetEvent).join(join_clause).where(where_clause),
+                order_by=[AssetEvent.timestamp],
+                session=session,
+            )
+
 
 def context_merge(context: Context, *args: Any, **kwargs: Any) -> None:
     """

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 (
-    AssetEventCollectionResponse,
+    AssetEventsResponse,
     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
-    ) -> AssetEventCollectionResponse:
+    ) -> AssetEventsResponse:
         """Get Asset event from the API server."""
         if name or uri:
             resp = self.client.get("asset-events/by-asset", params={"name": name, "uri": uri})
@@ -349,7 +349,7 @@ def get(
         else:
             raise ValueError("Either `name`, `uri` or `alias_name` must be provided")
 
-        return AssetEventCollectionResponse.model_validate_json(resp.read())
+        return AssetEventsResponse.model_validate_json(resp.read())
 
 
 class BearerAuth(httpx.Auth):

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

@@ -368,7 +368,7 @@ class AssetEventResponse(BaseModel):
     timestamp: Annotated[datetime, Field(title="Timestamp")]
 
 
-class AssetEventCollectionResponse(BaseModel):
+class AssetEventsResponse(BaseModel):
     """
     Collection of AssetEventResponse.
     """

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 (
-    AssetEventCollectionResponse,
+    AssetEventsResponse,
     AssetResponse,
     BundleInfo,
     ConnectionResponse,
@@ -105,27 +105,25 @@ def from_asset_response(cls, asset_response: AssetResponse) -> AssetResult:
         return cls(**asset_response.model_dump(exclude_defaults=True), type="AssetResult")
 
 
-class AssetEventCollectionResult(AssetEventCollectionResponse):
+class AssetEventsResult(AssetEventsResponse):
     """Response to GetAssetEvent request."""
 
-    type: Literal["AssetEventCollectionResult"] = "AssetEventCollectionResult"
+    type: Literal["AssetEventsResult"] = "AssetEventsResult"
 
     @classmethod
-    def from_asset_event_collection_response(
-        cls, asset_event_collection_response: AssetEventCollectionResponse
-    ) -> AssetEventCollectionResult:
+    def from_asset_events_response(cls, asset_events_response: AssetEventsResponse) -> AssetEventsResult:
         """
-        Get AssetEventCollectionResult from AssetEventCollectionResponse.
+        Get AssetEventsResult from AssetEventsResponse.
 
-        AssetEventCollectionResponse is autogenerated from the API schema, so we need to convert it to AssetEventCollectionResult
+        AssetEventsResponse is autogenerated from the API schema, so we need to convert it to AssetEventsResponse
         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_collection_response.model_dump(exclude_defaults=True),
-            type="AssetEventCollectionResult",
+            **asset_events_response.model_dump(exclude_defaults=True),
+            type="AssetEventsResult",
         )
 
 
@@ -209,7 +207,7 @@ class OKResponse(BaseModel):
 ToTask = Annotated[
     Union[
         AssetResult,
-        AssetEventCollectionResult,
+        AssetEventsResult,
         ConnectionResult,
         ErrorResponse,
         PrevSuccessfulDagRunResult,

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

@@ -48,6 +48,7 @@
     from airflow.sdk.definitions.context import Context
     from airflow.sdk.definitions.variable import Variable
     from airflow.sdk.execution_time.comms import (
+        AssetEventsResult,
         AssetResult,
         ConnectionResult,
         PrevSuccessfulDagRunResponse,
@@ -313,13 +314,6 @@ def __len__(self) -> int:
 
     def __getitem__(self, key: int | Asset | AssetAlias | AssetRef):
         from airflow.sdk.definitions.asset import Asset
-        from airflow.sdk.execution_time.comms import (
-            AssetEventCollectionResult,
-            ErrorResponse,
-            GetAssetEventByAsset,
-            GetAssetEventByAssetAlias,
-        )
-        from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
 
         if isinstance(key, int):  # Support index access; it's easier for trivial cases.
             obj = self._inlets[key]
@@ -328,6 +322,17 @@ def __getitem__(self, key: int | Asset | AssetAlias | AssetRef):
         else:
             obj = key
 
+        return self._get_asset_events_from_db(obj)
+
+    # TODO: This is temporary to avoid code duplication between here & airflow/models/taskinstance.py
+    def _get_asset_events_from_db(self, obj: Asset | AssetAlias | AssetRef) -> list[AssetEvent]:
+        from airflow.sdk.execution_time.comms import (
+            ErrorResponse,
+            GetAssetEventByAsset,
+            GetAssetEventByAssetAlias,
+        )
+        from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+
         if isinstance(obj, Asset):
             asset = self._assets[AssetUniqueKey.from_asset(obj)]
             SUPERVISOR_COMMS.send_request(log=log, msg=GetAssetEventByAsset(name=asset.name, uri=asset.uri))
@@ -346,15 +351,13 @@ def __getitem__(self, key: int | Asset | AssetAlias | AssetRef):
         elif isinstance(obj, AssetAlias):
             asset_alias = self._asset_aliases[AssetAliasUniqueKey.from_asset_alias(obj)]
             SUPERVISOR_COMMS.send_request(log=log, msg=GetAssetEventByAssetAlias(alias_name=asset_alias.name))
-        else:
-            raise ValueError(key)
 
         msg = SUPERVISOR_COMMS.get_message()
         if isinstance(msg, ErrorResponse):
             raise AirflowRuntimeError(msg)
 
         if TYPE_CHECKING:
-            assert isinstance(msg, AssetEventCollectionResult)
+            assert isinstance(msg, AssetEventsResult)
         return [AssetEvent(**event) for event in msg.model_dump()["asset_events"]]
 
 

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

@@ -61,7 +61,7 @@
     VariableResponse,
 )
 from airflow.sdk.execution_time.comms import (
-    AssetEventCollectionResult,
+    AssetEventsResult,
     AssetResult,
     ConnectionResult,
     DeferTask,
@@ -829,15 +829,11 @@ def _handle_request(self, msg: ToSupervisor, log: FilteringBoundLogger):
             resp = asset_result.model_dump_json(exclude_unset=True).encode()
         elif isinstance(msg, GetAssetEventByAsset):
             asset_event_resp = self.client.asset_events.get(uri=msg.uri, name=msg.name)
-            asset_event_result = AssetEventCollectionResult.from_asset_event_collection_response(
-                asset_event_resp
-            )
+            asset_event_result = AssetEventsResult.from_asset_events_response(asset_event_resp)
             resp = asset_event_result.model_dump_json(exclude_unset=True).encode()
         elif isinstance(msg, GetAssetEventByAssetAlias):
             asset_event_resp = self.client.asset_events.get(name=msg.alias_name)
-            asset_event_result = AssetEventCollectionResult.from_asset_event_collection_response(
-                asset_event_resp
-            )
+            asset_event_result = AssetEventsResult.from_asset_events_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)