From ff467e79e80ad9342b394fcbbe79e31003380890 Mon Sep 17 00:00:00 2001 From: arnoldmr01 Date: Thu, 20 Nov 2025 02:17:40 -0500 Subject: [PATCH] fix: mypy error in dagrun.py --- airflow-core/src/airflow/models/dagrun.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) 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.