pierrejeambrun commented on code in PR #71915:
URL: https://github.com/apache/airflow/pull/71915#discussion_r3872929132
##########
airflow-core/src/airflow/api_fastapi/core_api/services/ui/dependencies.py:
##########
@@ -599,15 +585,31 @@ def get_data_dependencies(
# Find other assets this entry task produces (outlets) to trace
downstream.
if task_key not in processed_tasks:
processed_tasks.add(task_key)
- outlet_refs = session.scalars(
- select(TaskOutletAssetReference).where(
- TaskOutletAssetReference.dag_id == dag_id,
- TaskOutletAssetReference.task_id == entry_task_id,
+ pending_outlet_keys.add(task_key)
+
+ if pending_inlet_keys:
+ inlet_refs = session.scalars(
+ select(TaskInletAssetReference).where(
+ tuple_(TaskInletAssetReference.dag_id,
TaskInletAssetReference.task_id).in_(
+ pending_inlet_keys
+ )
+ )
+ ).all()
+ for inlet_ref in inlet_refs:
+ if inlet_ref.asset_id not in processed_assets:
+ next_frontier.add(inlet_ref.asset_id)
+
+ if pending_outlet_keys:
+ outlet_refs = session.scalars(
+ select(TaskOutletAssetReference).where(
+ tuple_(TaskOutletAssetReference.dag_id,
TaskOutletAssetReference.task_id).in_(
+ pending_outlet_keys
)
- ).all()
- for outlet_ref in outlet_refs:
- if outlet_ref.asset_id not in processed_assets:
- next_frontier.add(outlet_ref.asset_id)
+ )
+ ).all()
+ for outlet_ref in outlet_refs:
+ if outlet_ref.asset_id not in processed_assets:
+ next_frontier.add(outlet_ref.asset_id)
Review Comment:
Pre-existing (not introduced here), but worth flagging: `processed_tasks` is
shared across the producer and consumer branches, so a task that both produces
one asset and consumes another (T reads X, writes Y) has whichever branch runs
first mark it processed — the other branch's lookup is silently skipped, and
any inlet/outlet unique to that side is dropped from the graph.
Not a blocker for this PR; possibly worth a follow-up to split into
`processed_inlet_tasks` / `processed_outlet_tasks`.
##########
airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_dependencies.py:
##########
@@ -868,6 +868,40 @@ def
test_data_dependencies_batches_entry_point_resolution_across_scheduled_dags(
entry_task_node_id = f"task:{dag_id}__SEPARATOR__entry_task"
assert nodes_by_id[entry_task_node_id]["type"] == "task"
+ def
test_data_dependencies_batches_inlet_outlet_resolution_across_producing_tasks(
Review Comment:
Test only operates on 'inlets' can we update it to also work on 'outlets'.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
