fix: fix mypy type error in dagrun.py

Author: arnoldmr01

airflow-core/src/airflow/models/dagrun.py

@@ -56,6 +56,7 @@
 from sqlalchemy.sql.expression import false, select
 from sqlalchemy.sql.functions import coalesce
 from sqlalchemy_utils import UUIDType
+from sqlalchemy.engine import ScalarResult
 
 from airflow.callbacks.callback_requests import DagCallbackRequest, DagRunContext
 from airflow.configuration import conf as airflow_conf
@@ -351,7 +352,7 @@ def __init__(
         if queued_at is NOTSET:
             self.queued_at = pendulum.now(tz="UTC") if state == DagRunState.QUEUED else None
         elif queued_at is not None:
-            self.queued_at = queued_at
+            self.queued_at = cast(datetime, queued_at)
         if run_type is not None:
             self.run_type = run_type
         self.creating_job_id = creating_job_id
@@ -572,11 +573,12 @@ def active_runs_of_dags(
         )
         if exclude_backfill:
             query = query.where(cls.run_type != DagRunType.BACKFILL_JOB)
-        return dict(session.execute(query).all())
+        result = session.execute(query).all()
+        return {dag_id:count for dag_id, count in result}
 
     @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 +617,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.
 
@@ -732,11 +734,13 @@ def find(
             qry = qry.where(cls.dag_id.in_(dag_ids))
 
         if is_container(run_id):
-            qry = qry.where(cls.run_id.in_(run_id))
+            run_ids = cast(Iterable[str], run_id)
+            qry = qry.where(cls.run_id.in_(run_ids))
         elif run_id is not None:
             qry = qry.where(cls.run_id == run_id)
         if is_container(logical_date):
-            qry = qry.where(cls.logical_date.in_(logical_date))
+            logical_dates = cast(Iterable[datetime], logical_date)
+            qry = qry.where(cls.logical_date.in_(logical_dates))
         elif logical_date is not None:
             qry = qry.where(cls.logical_date == logical_date)
         if logical_start_date and logical_end_date: