Store and expose task inlet references to assets (apache#51424)

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

@@ -0,0 +1,17 @@
+The ``consuming_dags`` key in asset API has been renamed to ``scheduled_dags``.
+
+The previous name caused confusion to users since the list does not contain all
+dags that technically *use* the asset, but only those that use it in their
+``schedule`` argument. As a bug fix, the key has been renamed to clarify its
+intention.
+
+* Types of change
+
+  * [ ] Dag changes
+  * [ ] Config changes
+  * [x] API changes
+  * [ ] CLI changes
+  * [ ] Behaviour changes
+  * [ ] Plugin changes
+  * [ ] Dependency changes
+  * [ ] Code interface changes

airflow-core/newsfragments/51424.significant.rst

@@ -46,6 +46,7 @@
     AssetAliasModel,
     AssetModel,
     DagScheduleAssetReference,
+    TaskInletAssetReference,
     TaskOutletAssetReference,
 )
 from airflow.models.connection import Connection
@@ -431,7 +432,7 @@ class _DagIdAssetReferenceFilter(BaseParam[list[str]]):
     """Search on dag_id."""
 
     def __init__(self, skip_none: bool = True) -> None:
-        super().__init__(AssetModel.consuming_dags, skip_none)
+        super().__init__(AssetModel.scheduled_dags, skip_none)
 
     @classmethod
     def depends(cls, dag_ids: list[str] = Query(None)) -> _DagIdAssetReferenceFilter:
@@ -444,8 +445,9 @@ 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.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."""
 
@@ -59,8 +68,9 @@ class AssetResponse(BaseModel):
     extra: dict | None = None
     created_at: datetime
     updated_at: datetime
-    consuming_dags: list[DagScheduleAssetReference]
+    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

@@ -7229,16 +7229,21 @@ components:
           type: string
           format: date-time
           title: Updated At
-        consuming_dags:
+        scheduled_dags:
           items:
             $ref: '#/components/schemas/DagScheduleAssetReference'
           type: array
-          title: Consuming Dags
+          title: Scheduled Dags
         producing_tasks:
           items:
             $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'
@@ -7256,8 +7261,9 @@ components:
       - group
       - created_at
       - updated_at
-      - consuming_dags
+      - 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

@@ -179,8 +179,9 @@ def get_assets(
 
     assets_rows = session.execute(
         assets_select.options(
-            subqueryload(AssetModel.consuming_dags),
+            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.consuming_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

@@ -130,7 +130,7 @@ def register_asset_change(
             .options(
                 joinedload(AssetModel.active),
                 joinedload(AssetModel.aliases),
-                joinedload(AssetModel.consuming_dags).joinedload(DagScheduleAssetReference.dag),
+                joinedload(AssetModel.scheduled_dags).joinedload(DagScheduleAssetReference.dag),
             )
         )
         if not asset_model:
@@ -161,7 +161,7 @@ def register_asset_change(
         session.flush()  # Ensure the event is written earlier than DDRQ entries below.
 
         dags_to_queue_from_asset = {
-            ref.dag for ref in asset_model.consuming_dags if not ref.dag.is_stale and not ref.dag.is_paused
+            ref.dag for ref in asset_model.scheduled_dags if not ref.dag.is_stale and not ref.dag.is_paused
         }
 
         dags_to_queue_from_asset_alias = set()
@@ -170,7 +170,7 @@ def register_asset_change(
                 select(AssetAliasModel)
                 .where(AssetAliasModel.name.in_(source_alias_names))
                 .options(
-                    joinedload(AssetAliasModel.consuming_dags).joinedload(DagScheduleAssetAliasReference.dag)
+                    joinedload(AssetAliasModel.scheduled_dags).joinedload(DagScheduleAssetAliasReference.dag)
                 )
             ).unique()
 
@@ -180,7 +180,7 @@ def register_asset_change(
 
                 dags_to_queue_from_asset_alias |= {
                     alias_ref.dag
-                    for alias_ref in asset_alias_model.consuming_dags
+                    for alias_ref in asset_alias_model.scheduled_dags
                     if not alias_ref.dag.is_stale and not alias_ref.dag.is_paused
                 }
 

airflow-core/src/airflow/assets/manager.py

@@ -44,6 +44,7 @@
     DagScheduleAssetNameReference,
     DagScheduleAssetReference,
     DagScheduleAssetUriReference,
+    TaskInletAssetReference,
     TaskOutletAssetReference,
 )
 from airflow.models.dag import DAG, DagModel, DagOwnerAttributes, DagTag
@@ -608,7 +609,7 @@ def _find_active_assets(name_uri_assets: Iterable[tuple[str, str]], session: Ses
             select(AssetModel.name, AssetModel.uri).where(
                 tuple_(AssetModel.name, AssetModel.uri).in_(name_uri_assets),
                 AssetModel.active.has(),
-                AssetModel.consuming_dags.any(
+                AssetModel.scheduled_dags.any(
                     DagScheduleAssetReference.dag.has(~DagModel.is_stale & ~DagModel.is_paused)
                 ),
             )
@@ -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:
@@ -949,7 +972,7 @@ def add_asset_trigger_references(
 
         # Remove references from assets no longer used
         orphan_assets = session.scalars(
-            select(AssetModel).filter(~AssetModel.consuming_dags.any()).filter(AssetModel.triggers.any())
+            select(AssetModel).filter(~AssetModel.scheduled_dags.any()).filter(AssetModel.triggers.any())
         )
         for asset_model in orphan_assets:
             if (asset_model.name, asset_model.uri) not in self.assets:

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

@@ -2375,7 +2375,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.consuming_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")