diff --git a/airflow-core/src/airflow/models/dagrun.py b/airflow-core/src/airflow/models/dagrun.py index 6727bdeb53c93..6651ca453a94e 100644 --- a/airflow-core/src/airflow/models/dagrun.py +++ b/airflow-core/src/airflow/models/dagrun.py @@ -98,7 +98,8 @@ from opentelemetry.sdk.trace import Span from pydantic import NonNegativeInt - from sqlalchemy.orm import Query, Session + from sqlalchemy.engine import ScalarResult + from sqlalchemy.orm import Session from sqlalchemy.sql.elements import Case, ColumnElement from airflow.models.dag_version import DagVersion @@ -572,11 +573,11 @@ def active_runs_of_dags( ) if exclude_backfill: query = query.where(cls.run_type != DagRunType.BACKFILL_JOB) - return dict(session.execute(query).all()) + return {dag_id: count for dag_id, count in session.execute(query)} @classmethod @retry_db_transaction - def get_running_dag_runs_to_examine(cls, session: Session) -> Query: + def get_running_dag_runs_to_examine(cls, session: Session) -> ScalarResult[DagRun]: """ Return the next DagRuns that the scheduler should attempt to schedule. @@ -615,7 +616,7 @@ def get_running_dag_runs_to_examine(cls, session: Session) -> Query: @classmethod @retry_db_transaction - def get_queued_dag_runs_to_set_running(cls, session: Session) -> Query: + def get_queued_dag_runs_to_set_running(cls, session: Session) -> ScalarResult[DagRun]: """ Return the next queued DagRuns that the scheduler should attempt to schedule.