def next_dagruns_to_examine( cls, state: DagRunState, session: Session, max_number: Optional[int] = None, ): """ Return the next DagRuns that the scheduler should attempt to schedule. This will return zero or more DagRun rows that are row-level-locked with a "SELECT ... FOR UPDATE" query, you should ensure that any scheduling decisions are made in a single transaction -- as soon as the transaction is committed it will be unlocked. :rtype: list[airflow.models.DagRun] """ from airflow.models.dag import DagModel if max_number is None: max_number = cls.DEFAULT_DAGRUNS_TO_EXAMINE # TODO: Bake this query, it is run _A lot_ query = ( session.query(cls) .filter(cls.state == state, cls.run_type != DagRunType.BACKFILL_JOB) .join( DagModel, DagModel.dag_id == cls.dag_id, ) .filter( DagModel.is_paused == expression.false(), DagModel.is_active == expression.true(), ) ) if state == State.QUEUED: # For dag runs in the queued state, we check if they have reached the max_active_runs limit # and if so we drop them running_drs = ( session.query(DagRun.dag_id, func.count(DagRun.state).label('num_running')) .filter(DagRun.state == DagRunState.RUNNING) .group_by(DagRun.dag_id) .subquery() ) query = query.outerjoin(running_drs, running_drs.c.dag_id == DagRun.dag_id).filter( func.coalesce(running_drs.c.num_running, 0) < DagModel.max_active_runs ) query = query.order_by( nulls_first(cls.last_scheduling_decision, session=session), cls.execution_date, ) if not settings.ALLOW_FUTURE_EXEC_DATES: query = query.filter(DagRun.execution_date <= func.now()) return with_row_locks( query.limit(max_number), of=cls, session=session, **skip_locked(session=session) )
def next_dagruns_to_examine( cls, state: DagRunState, session: Session, max_number: Optional[int] = None, ): """ Return the next DagRuns that the scheduler should attempt to schedule. This will return zero or more DagRun rows that are row-level-locked with a "SELECT ... FOR UPDATE" query, you should ensure that any scheduling decisions are made in a single transaction -- as soon as the transaction is committed it will be unlocked. :rtype: list[airflow.models.DagRun] """ from airflow.models.dag import DagModel if max_number is None: max_number = cls.DEFAULT_DAGRUNS_TO_EXAMINE # TODO: Bake this query, it is run _A lot_ query = ( session.query(cls) .filter(cls.state == state, cls.run_type != DagRunType.BACKFILL_JOB) .join( DagModel, DagModel.dag_id == cls.dag_id, ) .filter( DagModel.is_paused == expression.false(), DagModel.is_active == expression.true(), ) .order_by( nulls_first(cls.last_scheduling_decision, session=session), cls.execution_date, ) ) if not settings.ALLOW_FUTURE_EXEC_DATES: query = query.filter(DagRun.execution_date <= func.now()) return with_row_locks( query.limit(max_number), of=cls, session=session, **skip_locked(session=session) )