uranusjr commented on code in PR #64571:
URL: https://github.com/apache/airflow/pull/64571#discussion_r3109686192
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/ui/assets.py:
##########
@@ -115,6 +119,76 @@ def next_run_assets(
if not event.pop("queued", None):
event["lastUpdate"] = None
+ # For partitioned DAGs: enrich events with per-asset received/required
counts,
+ # using to_upstream for rollup mappers, and fix lastUpdate for partial
receipt.
+ if is_partitioned:
+ timetable = _load_timetable(dag_id, session)
+
+ pending_apdr = session.execute(
+ select(AssetPartitionDagRun.id, AssetPartitionDagRun.partition_key)
+ .where(
+ AssetPartitionDagRun.target_dag_id == dag_id,
+ AssetPartitionDagRun.created_dag_run_id.is_(None),
+ )
+ .order_by(AssetPartitionDagRun.created_at.desc())
+ .limit(1)
+ ).one_or_none()
+
+ if pending_apdr is not None:
+ # Collect received upstream partition keys per asset for this
partition run.
+ # Use a set to deduplicate: multiple events for the same key count
as one.
+ received_keys_by_asset: dict[int, set[str]] = {}
+ for row in session.execute(
+ select(
+ PartitionedAssetKeyLog.asset_id,
+ PartitionedAssetKeyLog.source_partition_key,
+ ).where(PartitionedAssetKeyLog.asset_partition_dag_run_id ==
pending_apdr.id)
+ ):
+ received_keys_by_asset.setdefault(row.asset_id,
set()).add(row.source_partition_key or "")
+
+ for event in events:
+ asset_id = event["id"]
+ received_keys = sorted(received_keys_by_asset.get(asset_id,
set()))
+ required_keys: list[str] = [pending_apdr.partition_key]
+ is_rollup = False
+ if timetable is not None:
+ with suppress(Exception):
+ mapper = timetable.get_partition_mapper(
+ name=event["name"],
+ uri=event["uri"],
+ )
+ if isinstance(mapper, RollupMapper):
+ required_keys =
sorted(mapper.to_upstream(pending_apdr.partition_key))
Review Comment:
I don’t like using `isinstance` here. This should be introspected on
properties.
--
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]
