nailo2c commented on code in PR #58543:
URL: https://github.com/apache/airflow/pull/58543#discussion_r3613447907
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -2551,6 +2568,128 @@ 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
+ triggered_date = timezone.coerce_datetime(max(record.created_at for
record in records))
+ asset_events = self._select_consumed_asset_events(
+ dag_id=dag.dag_id,
+ catchup=dag.catchup,
+ triggered_date=triggered_date,
+ previous_run_type=DagRunType.SCHEDULED,
+ session=session,
+ )
+ 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,
+ skip_locked=True,
+ key_share=False,
+ session=session,
+ )
+ ).all()
+
+ def _select_consumed_asset_events(
+ self,
+ *,
+ dag_id: str,
+ catchup: bool,
+ triggered_date: DateTime,
+ previous_run_type: DagRunType,
+ session: Session,
+ ) -> list[AssetEvent]:
+ """Select the asset events a new run consumes: events since the
previous run of the given type."""
+ previous_run_cte = (
+ select(func.max(DagRun.run_after).label("previous_run_after"))
+ .where(
+ DagRun.dag_id == dag_id,
+ DagRun.run_type == previous_run_type,
+ DagRun.run_after < triggered_date,
+ )
+ .cte()
+ )
+ # A first run has no previous run to floor the event window. With
catchup
+ # off, floor it at when the Dag started scheduling on its assets so the
+ # backlog is skipped; with catchup on, only date.min applies and the
+ # backlog replays.
+ event_window_floor: list[Any] = [previous_run_cte.c.previous_run_after]
+ if not catchup:
+ event_window_floor.append(
+ select(func.min(DagScheduleAssetReference.created_at))
+ .where(DagScheduleAssetReference.dag_id == dag_id)
+ .scalar_subquery()
+ )
+ event_window_floor.append(date.min)
+ return list(
+ session.scalars(
+ select(AssetEvent)
+ .where(
+ or_(
+ AssetEvent.asset_id.in_(
+ select(DagScheduleAssetReference.asset_id).where(
+ DagScheduleAssetReference.dag_id == dag_id
+ )
+ ),
+ AssetEvent.source_aliases.any(
+ AssetAliasModel.scheduled_dags.any(
+ DagScheduleAssetAliasReference.dag_id == dag_id
+ )
+ ),
+ ),
+ AssetEvent.timestamp <= triggered_date,
+ AssetEvent.timestamp > func.coalesce(*event_window_floor),
+ # A gated run consumes events newer than its run_after
(the slot
+ # time), so the previous-run floor alone would
re-attribute them to
+ # the next run; skip events already consumed by this Dag's
runs.
+ ~AssetEvent.created_dagruns.any(DagRun.dag_id == dag_id),
Review Comment:
ah, good catch, 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]
