aaron-y-chen commented on code in PR #58543:
URL: https://github.com/apache/airflow/pull/58543#discussion_r4108270685
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -2776,6 +2793,113 @@ def _create_dag_runs(self, dag_models:
Collection[DagModel], session: Session) -
# TODO[HA]: Should we do a session.flush() so we don't have to
keep lots of state/object in
# memory for larger dags? or expunge_all()
+ def _collect_gated_asset_events(
+ self, *, dag: SerializedDAG, session: Session
+ ) -> tuple[Sequence[AssetDagRunQueue], list[AssetEvent]] | None:
+ """
+ Check an asset-gated Dag's asset condition and collect what a new run
consumes.
+
+ Returns ``None`` when the condition is not satisfied by the queued
asset
+ events, in which case no run should be created yet.
+ """
+ records = self._lock_queued_asset_records(dag_id=dag.dag_id,
session=session)
+ if not records:
+ return None
+ statuses = {SerializedAssetUniqueKey.from_asset(record.asset): True
for record in records}
+ try:
+ ready = AssetEvaluator(session).run(dag.timetable.asset_condition,
statuses=statuses)
+ except Exception:
+ self.log.exception("Dag '%s' failed to be evaluated; assuming not
ready", dag.dag_id)
+ return None
+ if not ready:
+ return None
+ asset_events = self._select_consumed_asset_events(
+ dag=dag,
+ records=records,
+ session=session,
+ )
+ if not asset_events:
+ self._delete_consumed_asset_records(records=records,
dag_id=dag.dag_id, session=session)
+ return None
+ return records, asset_events
+
+ def _lock_queued_asset_records(self, *, dag_id: str, session: Session) ->
Sequence[AssetDagRunQueue]:
+ """Lock and return the Dag's queued asset (ADRQ) rows, skipping rows
another scheduler holds."""
+ return session.scalars(
+ with_row_locks(
+ select(AssetDagRunQueue)
+ .where(AssetDagRunQueue.target_dag_id == dag_id)
+ .options(joinedload(AssetDagRunQueue.asset)),
+ of=AssetDagRunQueue,
Review Comment:
Nice suggestion, I'd prefer the more efficient approach. Let me fix it.
--
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]
