Store inlet references in database and expose API

Author: uranusjr

airflow-core/docs/img/airflow_erd.sha256

@@ -1 +1 @@
-789370f7c29cc290c307562a13f06576d8651e12916cc42f89653f8127c00210
+7e97c71ee6da77d758087a5d7129ac8d63b52dcc3bf4c46b7faf5977faa53ef4

airflow-core/docs/img/airflow_erd.svg

@@ -39,7 +39,9 @@ Here's the list of all the Database Migrations that are executed via when you ru
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+
 | Revision ID             | Revises ID       | Airflow Version   | Description                                                  |
 +=========================+==================+===================+==============================================================+
-| ``3ac9e5732b1f`` (head) | ``0242ac120002`` | ``3.1.0``         | Change the on-delete behaviour of                            |
+| ``583e80dfcef4`` (head) | ``3ac9e5732b1f`` | ``3.1.0``         | Add task_inlet_asset_reference table.                        |
++-------------------------+------------------+-------------------+--------------------------------------------------------------+
+| ``3ac9e5732b1f``        | ``0242ac120002`` | ``3.1.0``         | Change the on-delete behaviour of                            |
 |                         |                  |                   | task_instance.dag_version_id foreign key constraint to       |
 |                         |                  |                   | RESTRICT.                                                    |
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+

airflow-core/docs/migrations-ref.rst

@@ -46,6 +46,7 @@
     AssetAliasModel,
     AssetModel,
     DagScheduleAssetReference,
+    TaskInletAssetReference,
     TaskOutletAssetReference,
 )
 from airflow.models.connection import Connection
@@ -446,6 +447,7 @@ def to_orm(self, select: Select) -> Select:
         return select.where(
             (AssetModel.scheduled_dags.any(DagScheduleAssetReference.dag_id.in_(self.value)))
             | (AssetModel.producing_tasks.any(TaskOutletAssetReference.dag_id.in_(self.value)))
+            | (AssetModel.consuming_tasks.any(TaskInletAssetReference.dag_id.in_(self.value)))
         )
 
 

airflow-core/src/airflow/api_fastapi/common/parameters.py

@@ -33,6 +33,15 @@ class DagScheduleAssetReference(StrictBaseModel):
     updated_at: datetime
 
 
+class TaskInletAssetReference(StrictBaseModel):
+    """Task inlet reference serializer for assets."""
+
+    dag_id: str
+    task_id: str
+    created_at: datetime
+    updated_at: datetime
+
+
 class TaskOutletAssetReference(StrictBaseModel):
     """Task outlet reference serializer for assets."""
 
@@ -61,6 +70,7 @@ class AssetResponse(BaseModel):
     updated_at: datetime
     scheduled_dags: list[DagScheduleAssetReference]
     producing_tasks: list[TaskOutletAssetReference]
+    consuming_tasks: list[TaskInletAssetReference]
     aliases: list[AssetAliasResponse]
     last_asset_event: LastAssetEventResponse | None = None
 

airflow-core/src/airflow/api_fastapi/core_api/datamodels/assets.py

@@ -7239,6 +7239,11 @@ components:
             $ref: '#/components/schemas/TaskOutletAssetReference'
           type: array
           title: Producing Tasks
+        consuming_tasks:
+          items:
+            $ref: '#/components/schemas/TaskInletAssetReference'
+          type: array
+          title: Consuming Tasks
         aliases:
           items:
             $ref: '#/components/schemas/AssetAliasResponse'
@@ -7258,6 +7263,7 @@ components:
       - updated_at
       - scheduled_dags
       - producing_tasks
+      - consuming_tasks
       - aliases
       title: AssetResponse
       description: Asset serializer for responses.
@@ -9999,6 +10005,31 @@ components:
       - reason
       title: TaskDependencyResponse
       description: Task Dependency serializer for responses.
+    TaskInletAssetReference:
+      properties:
+        dag_id:
+          type: string
+          title: Dag Id
+        task_id:
+          type: string
+          title: Task Id
+        created_at:
+          type: string
+          format: date-time
+          title: Created At
+        updated_at:
+          type: string
+          format: date-time
+          title: Updated At
+      additionalProperties: false
+      type: object
+      required:
+      - dag_id
+      - task_id
+      - created_at
+      - updated_at
+      title: TaskInletAssetReference
+      description: Task inlet reference serializer for assets.
     TaskInstanceCollectionResponse:
       properties:
         task_instances:

airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml

@@ -181,6 +181,7 @@ def get_assets(
         assets_select.options(
             subqueryload(AssetModel.scheduled_dags),
             subqueryload(AssetModel.producing_tasks),
+            subqueryload(AssetModel.consuming_tasks),
         )
     )
 
@@ -456,7 +457,11 @@ def get_asset(
     asset = session.scalar(
         select(AssetModel)
         .where(AssetModel.id == asset_id)
-        .options(joinedload(AssetModel.scheduled_dags), joinedload(AssetModel.producing_tasks))
+        .options(
+            joinedload(AssetModel.scheduled_dags),
+            joinedload(AssetModel.producing_tasks),
+            joinedload(AssetModel.consuming_tasks),
+        )
     )
 
     last_asset_event_id = asset_event_rows[1] if asset_event_rows else None

airflow-core/src/airflow/api_fastapi/core_api/routes/public/assets.py

@@ -44,6 +44,7 @@
     DagScheduleAssetNameReference,
     DagScheduleAssetReference,
     DagScheduleAssetUriReference,
+    TaskInletAssetReference,
     TaskOutletAssetReference,
 )
 from airflow.models.dag import DAG, DagModel, DagOwnerAttributes, DagTag
@@ -623,6 +624,7 @@ class AssetModelOperation(NamedTuple):
     schedule_asset_alias_references: dict[str, list[AssetAlias]]
     schedule_asset_name_references: set[tuple[str, str]]  # dag_id, ref_name.
     schedule_asset_uri_references: set[tuple[str, str]]  # dag_id, ref_uri.
+    inlet_references: dict[str, list[tuple[str, Asset]]]
     outlet_references: dict[str, list[tuple[str, Asset]]]
     assets: dict[tuple[str, str], Asset]
     asset_aliases: dict[str, AssetAlias]
@@ -650,6 +652,9 @@ def collect(cls, dags: dict[str, MaybeSerializedDAG]) -> Self:
                 for ref in dag.timetable.asset_condition.iter_asset_refs()
                 if isinstance(ref, AssetUriRef)
             },
+            inlet_references={
+                dag_id: list(dag.get_task_assets(inlets=True, outlets=False)) for dag_id, dag in dags.items()
+            },
             outlet_references={
                 dag_id: list(dag.get_task_assets(inlets=False, outlets=True)) for dag_id, dag in dags.items()
             },
@@ -829,6 +834,24 @@ def add_task_asset_references(
         # Optimization: No assets means there are no references to update.
         if not assets:
             return
+        for dag_id, references in self.inlet_references.items():
+            # Optimization: no references at all; this is faster than repeated delete().
+            if not references:
+                dags[dag_id].task_inlet_asset_references = []
+                continue
+            referenced_inlets = {
+                (task_id, asset.id)
+                for task_id, asset in ((task_id, assets[d.name, d.uri]) for task_id, d in references)
+            }
+            orm_refs = {(r.task_id, r.asset_id): r for r in dags[dag_id].task_inlet_asset_references}
+            for key, ref in orm_refs.items():
+                if key not in referenced_inlets:
+                    session.delete(ref)
+            session.bulk_save_objects(
+                TaskInletAssetReference(asset_id=asset_id, dag_id=dag_id, task_id=task_id)
+                for task_id, asset_id in referenced_inlets
+                if (task_id, asset_id) not in orm_refs
+            )
         for dag_id, references in self.outlet_references.items():
             # Optimization: no references at all; this is faster than repeated delete().
             if not references:

airflow-core/src/airflow/dag_processing/collection.py

@@ -2366,7 +2366,12 @@ def _activate_referenced_assets(assets: Collection[AssetModel], *, session: Sess
         def _generate_warning_message(
             offending: AssetModel, attr: str, value: str
         ) -> Iterator[tuple[str, str]]:
-            for ref in itertools.chain(offending.scheduled_dags, offending.producing_tasks):
+            offending_references = itertools.chain(
+                offending.scheduled_dags,
+                offending.producing_tasks,
+                offending.consuming_tasks,
+            )
+            for ref in offending_references:
                 yield (
                     ref.dag_id,
                     (

airflow-core/src/airflow/jobs/scheduler_job_runner.py

@@ -0,0 +1,70 @@
+#
+# 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.
+
+"""
+Add task_inlet_asset_reference table.
+
+Revision ID: 583e80dfcef4
+Revises: 3ac9e5732b1f
+Create Date: 2025-06-04 06:26:36.536172
+"""
+
+from __future__ import annotations
+
+from alembic import op
+from sqlalchemy import Column, ForeignKeyConstraint, Index, Integer, PrimaryKeyConstraint
+
+from airflow.migrations.db_types import StringID
+from airflow.utils.sqlalchemy import UtcDateTime
+
+revision = "583e80dfcef4"
+down_revision = "3ac9e5732b1f"
+branch_labels = None
+depends_on = None
+airflow_version = "3.1.0"
+
+
+def upgrade():
+    """Add task_inlet_asset_reference table."""
+    op.create_table(
+        "task_inlet_asset_reference",
+        Column("asset_id", Integer, primary_key=True, nullable=False),
+        Column("dag_id", StringID(), primary_key=True, nullable=False),
+        Column("task_id", StringID(), primary_key=True, nullable=False),
+        Column("created_at", UtcDateTime, nullable=False),
+        Column("updated_at", UtcDateTime, nullable=False),
+        ForeignKeyConstraint(
+            ["asset_id"],
+            ["asset.id"],
+            name="tiar_asset_fkey",
+            ondelete="CASCADE",
+        ),
+        PrimaryKeyConstraint("asset_id", "dag_id", "task_id", name="tiar_pkey"),
+        ForeignKeyConstraint(
+            columns=["dag_id"],
+            refcolumns=["dag.dag_id"],
+            name="tiar_dag_id_fkey",
+            ondelete="CASCADE",
+        ),
+        Index("idx_task_inlet_asset_reference_dag_id", "dag_id"),
+    )
+
+
+def downgrade():
+    """Remove task_inlet_asset_reference table."""
+    op.drop_table("task_inlet_asset_reference")

airflow-core/src/airflow/migrations/versions/0073_3_1_0_add_task_inlet_asset_reference.py

@@ -274,6 +274,7 @@ class AssetModel(Base):
 
     scheduled_dags = relationship("DagScheduleAssetReference", back_populates="asset")
     producing_tasks = relationship("TaskOutletAssetReference", back_populates="asset")
+    consuming_tasks = relationship("TaskInletAssetReference", back_populates="asset")
     triggers = relationship("Trigger", secondary=asset_trigger_association_table, back_populates="assets")
 
     __tablename__ = "asset"
@@ -612,6 +613,50 @@ def __repr__(self):
         return f"{self.__class__.__name__}({', '.join(args)})"
 
 
+class TaskInletAssetReference(Base):
+    """References from a task to an asset that it references as an inlet."""
+
+    asset_id = Column(Integer, primary_key=True, nullable=False)
+    dag_id = Column(StringID(), primary_key=True, nullable=False)
+    task_id = Column(StringID(), primary_key=True, nullable=False)
+    created_at = Column(UtcDateTime, default=timezone.utcnow, nullable=False)
+    updated_at = Column(UtcDateTime, default=timezone.utcnow, onupdate=timezone.utcnow, nullable=False)
+
+    asset = relationship("AssetModel", back_populates="consuming_tasks")
+
+    __tablename__ = "task_inlet_asset_reference"
+    __table_args__ = (
+        ForeignKeyConstraint(
+            (asset_id,),
+            ["asset.id"],
+            name="tiar_asset_fkey",
+            ondelete="CASCADE",
+        ),
+        PrimaryKeyConstraint(asset_id, dag_id, task_id, name="tiar_pkey"),
+        ForeignKeyConstraint(
+            columns=(dag_id,),
+            refcolumns=["dag.dag_id"],
+            name="tiar_dag_id_fkey",
+            ondelete="CASCADE",
+        ),
+        Index("idx_task_inlet_asset_reference_dag_id", dag_id),
+    )
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, self.__class__):
+            return NotImplemented
+        return (
+            self.asset_id == other.asset_id and self.dag_id == other.dag_id and self.task_id == other.task_id
+        )
+
+    def __hash__(self):
+        return hash(self.__mapper__.primary_key)
+
+    def __repr__(self):
+        args = (f"{(attr := x.name)}={getattr(self, attr)!r}" for x in self.__mapper__.primary_key)
+        return f"{self.__class__.__name__}({', '.join(args)})"
+
+
 class AssetDagRunQueue(Base):
     """Model for storing asset events that need processing."""
 

airflow-core/src/airflow/models/asset.py

@@ -1756,7 +1756,7 @@ def get_task_assets(
         of_type: type[AssetT] = Asset,  # type: ignore[assignment]
     ) -> Generator[tuple[str, AssetT], None, None]:
         for task in self.task_dict.values():
-            directions = ("inlets",) if inlets else ()
+            directions: tuple[str, ...] = ("inlets",) if inlets else ()
             if outlets:
                 directions += ("outlets",)
             for direction in directions:
@@ -1966,6 +1966,10 @@ class DagModel(Base):
         cascade="all, delete, delete-orphan",
     )
     schedule_assets = association_proxy("schedule_asset_references", "asset")
+    task_inlet_asset_references = relationship(
+        "TaskInletAssetReference",
+        cascade="all, delete, delete-orphan",
+    )
     task_outlet_asset_references = relationship(
         "TaskOutletAssetReference",
         cascade="all, delete, delete-orphan",