provide_session keep return type (#9787)

GitOrigin-RevId: 514eb6d1e350818b31dca5adeaec2d7fd32b23ee

airflow/api_connexion/schemas/pool_schema.py

@@ -59,7 +59,7 @@ def get_queued_slots(obj: Pool) -> int:
         return obj.queued_slots()
 
     @staticmethod
-    def get_open_slots(obj: Pool) -> int:
+    def get_open_slots(obj: Pool) -> float:
         """
         Returns the open slots of the pool.
         """

airflow/jobs/backfill_job.py

@@ -20,7 +20,7 @@
 import time
 from collections import OrderedDict
 from datetime import datetime
-from typing import Set
+from typing import Optional, Set
 
 from sqlalchemy.orm.session import Session, make_transient
 from tabulate import tabulate
@@ -295,12 +295,14 @@ def _get_dag_run(self, run_date: datetime, dag: DAG, session: Session = None):
 
         # check if we are scheduling on top of a already existing dag_run
         # we could find a "scheduled" run instead of a "backfill"
-        run = DagRun.find(dag_id=dag.dag_id,
-                          execution_date=run_date,
-                          session=session)
-
-        if run is not None and len(run) > 0:
-            run = run[0]
+        runs = DagRun.find(
+            dag_id=dag.dag_id,
+            execution_date=run_date,
+            session=session
+        )
+        run: Optional[DagRun]
+        if runs:
+            run = runs[0]
             if run.state == State.RUNNING:
                 respect_dag_max_active_limit = False
         else:

airflow/utils/session.py

@@ -17,6 +17,7 @@
 
 import contextlib
 from functools import wraps
+from typing import Callable, TypeVar
 
 from airflow import settings
 
@@ -37,15 +38,18 @@ def create_session():
         session.close()
 
 
-def provide_session(func):
+RT = TypeVar("RT")  # pylint: disable=invalid-name
+
+
+def provide_session(func: Callable[..., RT]) -> Callable[..., RT]:
     """
     Function decorator that provides a session if it isn't provided.
     If you want to reuse a session or run the function as part of a
     database transaction, you pass it to the function, if not this wrapper
     will create one and close it for you.
     """
     @wraps(func)
-    def wrapper(*args, **kwargs):
+    def wrapper(*args, **kwargs) -> RT:
         arg_session = 'session'
 
         func_params = func.__code__.co_varnames