Repository: apache/airflow
File: airflow-core/src/airflow/jobs/scheduler_job_runner.py (Python)

Code before:
    def load(self, session: Session) -> None:
        self.dag_run_active_tasks_map.clear()
        self.task_concurrency_map.clear()
        self.task_dagrun_concurrency_map.clear()
        query = session.execute(
            select(TI.dag_id, TI.task_id, TI.run_id, func.count("*"))
            .where(TI.state.in_(EXECUTION_STATES))
            .group_by(TI.task_id, TI.run_id, TI.dag_id)
        )
        for dag_id, task_id, run_id, c in query:
            self.dag_run_active_tasks_map[dag_id, run_id] += c
            self.task_concurrency_map[(dag_id, task_id)] += c
            self.task_dagrun_concurrency_map[(dag_id, run_id, task_id)] += c


def _is_parent_process() -> bool:
    """
    Whether this is a parent process.

    Return True if the current process is the parent process.
    False if the current process is a child process started by multiprocessing.
    """
    return multiprocessing.current_process().name == "MainProcess"


def _get_current_dr_task_concurrency(states: Iterable[TaskInstanceState]) -> Query:
    """Get the dag_run IDs and how many tasks are in the provided states for each one."""
    return (
        select(TI.run_id, func.count("*").label("dr_count"))
        .where(TI.state.in_(states))
        .group_by(TI.run_id)
        .subquery()
    )


class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
    """
    SchedulerJobRunner runs for a specific time interval and schedules jobs that are ready to run.

    It figures out the latest runs for each task and sees if the dependencies
    for the next schedules are met.
    If so, it creates appropriate TaskInstances and sends run commands to the
    executor. It does this for each task in each DAG and repeats.

    :param num_runs: The number of times to run the scheduling loop. If you
        have a large number of DAG files this could complete before each file
        has been parsed. -1 for unlimited times.
    :param scheduler_idle_sleep_time: The number of seconds to wait between
        polls of running processors
    :param log: override the default Logger
    """

Review comment (suggestion):
I _think_ this return value is a Subquery, not a Query (honestly the return type annotation is pretty useless anyway due to the lack of SQLAlchemy support)

Code after:

def _eager_load_dag_run_for_validation() -> tuple[LoaderOption, LoaderOption]:
    """
    Eager-load DagRun relations required for execution API datamodel validation.

    When building TaskCallbackRequest with DRDataModel.model_validate(ti.dag_run),
    the consumed_asset_events collection and nested asset/source_aliases must be
    preloaded to avoid DetachedInstanceError after the session closes.

    Returns a tuple of two load options:
        - Asset loader: TI.dag_run → consumed_asset_events → asset
        - Alias loader: TI.dag_run → consumed_asset_events → source_aliases

    Example usage::

        asset_loader, alias_loader = _eager_load_dag_run_for_validation()
        query = select(TI).options(asset_loader).options(alias_loader)
    """
    # Traverse TI → dag_run → consumed_asset_events once, then branch to asset/aliases
    base = joinedload(TI.dag_run).selectinload(DagRun.consumed_asset_events)
    return (
        base.selectinload(AssetEvent.asset),
        base.selectinload(AssetEvent.source_aliases),
    )


def _get_current_dag(dag_id: str, session: Session) -> SerializedDAG | None:
    serdag = SerializedDagModel.get(dag_id=dag_id, session=session)  # grabs the latest version
    if not serdag:
        return None
    serdag.load_op_links = False
    return serdag.dag


class ConcurrencyMap:
    """
    Dataclass to represent concurrency maps.

    It contains a map from (dag_id, task_id) to # of task instances, a map from (dag_id, task_id)
    to # of task instances in the given state list and a map from (dag_id, run_id, task_id)
    to # of task instances in the given state list in each DAG run.
    """

    def __init__(self):
        self.dag_run_active_tasks_map: Counter[tuple[str, str]] = Counter()
        self.task_concurrency_map: Counter[tuple[str, str]] = Counter()
        self.task_dagrun_concurrency_map: Counter[tuple[str, str, str]] = Counter()

    def load(self, session: Session) -> None:
        self.dag_run_active_tasks_map.clear()
        self.task_concurrency_map.clear()