ColtenOuO commented on code in PR #71915:
URL: https://github.com/apache/airflow/pull/71915#discussion_r3887751191
##########
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:
Thanks! I've looked into this and tried to address it in PR #72272.
Really appreciate that you always provide great follow-up ideas and
potential improvements when reviewing PRs!
--
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]
