feat(task_sdk): add support for inlet_events in Task Context

Author: Lee-W

airflow/api_fastapi/execution_api/datamodels/asset.py

@@ -17,7 +17,12 @@
 
 from __future__ import annotations
 
+from datetime import datetime
+
+from pydantic import Field, field_validator
+
 from airflow.api_fastapi.core_api.base import BaseModel, StrictBaseModel
+from airflow.sdk.execution_time.secrets_masker import redact
 
 
 class AssetResponse(BaseModel):
@@ -36,7 +41,42 @@ class AssetAliasResponse(BaseModel):
     group: str
 
 
-class AssetProfile(StrictBaseModel):
+class DagRunAssetReference(StrictBaseModel):
+    """DAGRun serializer for asset responses."""
+
+    run_id: str
+    dag_id: str
+    execution_date: datetime = Field(alias="logical_date")
+    start_date: datetime
+    end_date: datetime | None
+    state: str
+    data_interval_start: datetime
+    data_interval_end: datetime
+
+
+class AssetEventResponse(BaseModel):
+    """Asset event schema with fields that are needed for Runtime."""
+
+    id: int
+    asset_id: int
+    uri: str | None = Field(alias="uri", default=None)
+    name: str | None = Field(alias="name", default=None)
+    group: str | None = Field(alias="group", default=None)
+    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
+
+    @field_validator("extra", mode="after")
+    @classmethod
+    def redact_extra(cls, v: dict):
+        return redact(v)
+
+
+class AssetProfile(BaseModel):
     """
     Profile of an Asset.
 

airflow/api_fastapi/execution_api/routes/__init__.py

@@ -18,6 +18,7 @@
 
 from airflow.api_fastapi.common.router import AirflowRouter
 from airflow.api_fastapi.execution_api.routes import (
+    asset_events,
     assets,
     connections,
     health,
@@ -28,6 +29,7 @@
 
 execution_api_router = AirflowRouter()
 execution_api_router.include_router(assets.router, prefix="/assets", tags=["Assets"])
+execution_api_router.include_router(asset_events.router, prefix="/asset-events", tags=["Asset Events"])
 execution_api_router.include_router(connections.router, prefix="/connections", tags=["Connections"])
 execution_api_router.include_router(health.router, tags=["Health"])
 execution_api_router.include_router(task_instances.router, prefix="/task-instances", tags=["Task Instances"])

airflow/api_fastapi/execution_api/routes/asset_events.py

@@ -0,0 +1,129 @@
+# 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 typing import TYPE_CHECKING, Annotated
+
+from fastapi import HTTPException, 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.models.asset import AssetAliasModel, AssetEvent, AssetModel
+from airflow.utils.db import LazySelectSequence
+
+if TYPE_CHECKING:
+    from sqlalchemy.engine import Row
+    from sqlalchemy.sql.expression import Select, TextClause
+
+# TODO: Add dependency on JWT token
+router = AirflowRouter(
+    responses={
+        status.HTTP_404_NOT_FOUND: {"description": "Asset not found"},
+        status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
+    },
+)
+
+
+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(
+    *, 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,
+    )
+    _raise_if_not_found(asset_event=asset_event, msg="Not found")
+    return AssetEventResponse.model_validate(asset_event)
+
+
+@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(
+        join_clause=AssetEvent.asset,
+        where_clause=and_(AssetModel.name == name, AssetModel.uri == uri),
+        session=session,
+    )
+
+
+@router.get("/by-asset-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(
+        join_clause=AssetEvent.asset,
+        where_clause=and_(AssetModel.uri == uri, AssetModel.active.has()),
+        session=session,
+    )
+
+
+@router.get("/by-asset-name")
+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(
+        join_clause=AssetEvent.asset,
+        where_clause=and_(AssetModel.uri == name, AssetModel.active.has()),
+        session=session,
+    )
+
+
+@router.get("/by-alias-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(
+        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,
+            },
+        )

airflow/utils/context.py

@@ -21,50 +21,33 @@
 
 from collections.abc import (
     Container,
-    Iterator,
-    Mapping,
 )
 from typing import (
     TYPE_CHECKING,
     Any,
-    Union,
     cast,
 )
 
 import attrs
-from sqlalchemy import and_, select
+from sqlalchemy import select
 
 from airflow.models.asset import (
-    AssetAliasModel,
-    AssetEvent,
     AssetModel,
-    fetch_active_assets_by_name,
-    fetch_active_assets_by_uri,
 )
 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.orm import Session
-    from sqlalchemy.sql.expression import Select, TextClause
-
     from airflow.sdk.types import OutletEventAccessorsProtocol
 
 # NOTE: Please keep this in sync with the following:
@@ -170,105 +153,14 @@ 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(Mapping[Union[int, Asset, AssetAlias, AssetRef], LazyAssetEventSelectSequence]):
+class InletEventsAccessors(InletEventsAccessorsSDK):
     """
     Lazy mapping for inlet asset events accessors.
 
     :meta private:
     """
 
-    _inlets: list[Any]
-    _assets: dict[AssetUniqueKey, Asset]
-    _asset_aliases: dict[AssetAliasUniqueKey, AssetAlias]
-    _session: Session
-
-    def __init__(self, inlets: list, *, session: Session) -> None:
-        self._inlets = inlets
-        self._session = session
-        self._assets = {}
-        self._asset_aliases = {}
-
-        _asset_ref_names: list[str] = []
-        _asset_ref_uris: list[str] = []
-        for inlet in inlets:
-            if isinstance(inlet, Asset):
-                self._assets[AssetUniqueKey.from_asset(inlet)] = inlet
-            elif isinstance(inlet, AssetAlias):
-                self._asset_aliases[AssetAliasUniqueKey.from_asset_alias(inlet)] = inlet
-            elif isinstance(inlet, AssetNameRef):
-                _asset_ref_names.append(inlet.name)
-            elif isinstance(inlet, AssetUriRef):
-                _asset_ref_uris.append(inlet.uri)
-
-        if _asset_ref_names:
-            for _, asset in fetch_active_assets_by_name(_asset_ref_names, self._session).items():
-                self._assets[AssetUniqueKey.from_asset(asset)] = asset
-        if _asset_ref_uris:
-            for _, asset in fetch_active_assets_by_uri(_asset_ref_uris, self._session).items():
-                self._assets[AssetUniqueKey.from_asset(asset)] = asset
-
-    def __iter__(self) -> Iterator[Asset | AssetAlias]:
-        return iter(self._inlets)
-
-    def __len__(self) -> int:
-        return len(self._inlets)
-
-    def __getitem__(self, key: int | Asset | AssetAlias | AssetRef) -> LazyAssetEventSelectSequence:
-        if isinstance(key, int):  # Support index access; it's easier for trivial cases.
-            obj = self._inlets[key]
-            if not isinstance(obj, (Asset, AssetAlias, AssetRef)):
-                raise IndexError(key)
-        else:
-            obj = key
-
-        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())
-        else:
-            raise ValueError(key)
-
-        return LazyAssetEventSelectSequence.from_select(
-            select(AssetEvent).join(join_clause).where(where_clause),
-            order_by=[AssetEvent.timestamp],
-            session=self._session,
-        )
-
 
 def context_merge(context: Context, *args: Any, **kwargs: Any) -> None:
     """