Refactor Sqlalchemy queries to 2.0 style (Part 5) (#32474)

Author: phanikumv

airflow/cli/commands/dag_command.py

@@ -28,7 +28,7 @@
 import warnings
 
 from graphviz.dot import Dot
-from sqlalchemy import delete
+from sqlalchemy import delete, select
 from sqlalchemy.orm import Session
 
 from airflow import settings
@@ -287,7 +287,7 @@ def dag_state(args, session: Session = NEW_SESSION) -> None:
 
     if not dag:
         raise SystemExit(f"DAG: {args.dag_id} does not exist in 'dag' table")
-    dr = session.query(DagRun).filter_by(dag_id=args.dag_id, execution_date=args.execution_date).one_or_none()
+    dr = session.scalar(select(DagRun).filter_by(dag_id=args.dag_id, execution_date=args.execution_date))
     out = dr.state if dr else None
     conf_out = ""
     if out and dr.conf:
@@ -309,7 +309,9 @@ def dag_next_execution(args) -> None:
         print("[INFO] Please be reminded this DAG is PAUSED now.", file=sys.stderr)
 
     with create_session() as session:
-        last_parsed_dag: DagModel = session.query(DagModel).filter(DagModel.dag_id == dag.dag_id).one()
+        last_parsed_dag: DagModel = session.scalars(
+            select(DagModel).where(DagModel.dag_id == dag.dag_id)
+        ).one()
 
     def print_execution_interval(interval: DataInterval | None):
         if interval is None:
@@ -428,8 +430,10 @@ def dag_list_jobs(args, dag: DAG | None = None, session: Session = NEW_SESSION)
         queries.append(Job.state == args.state)
 
     fields = ["dag_id", "state", "job_type", "start_date", "end_date"]
-    all_jobs = session.query(Job).filter(*queries).order_by(Job.start_date.desc()).limit(args.limit).all()
-    all_jobs = [{f: str(job.__getattribute__(f)) for f in fields} for job in all_jobs]
+    all_jobs_iter = session.scalars(
+        select(Job).where(*queries).order_by(Job.start_date.desc()).limit(args.limit)
+    )
+    all_jobs = [{f: str(job.__getattribute__(f)) for f in fields} for job in all_jobs_iter]
 
     AirflowConsole().print_as(
         data=all_jobs,
@@ -492,14 +496,12 @@ def dag_test(args, dag: DAG | None = None, session: Session = NEW_SESSION) -> No
     imgcat = args.imgcat_dagrun
     filename = args.save_dagrun
     if show_dagrun or imgcat or filename:
-        tis = (
-            session.query(TaskInstance)
-            .filter(
+        tis = session.scalars(
+            select(TaskInstance).where(
                 TaskInstance.dag_id == args.dag_id,
                 TaskInstance.execution_date == execution_date,
             )
-            .all()
-        )
+        ).all()
 
         dot_graph = render_dag(dag, tis=tis)
         print()

airflow/cli/commands/jobs_command.py

@@ -16,6 +16,7 @@
 # under the License.
 from __future__ import annotations
 
+from sqlalchemy import select
 from sqlalchemy.orm import Session
 
 from airflow.jobs.job import Job
@@ -32,17 +33,17 @@ def check(args, session: Session = NEW_SESSION) -> None:
     if args.hostname and args.local:
         raise SystemExit("You can't use --hostname and --local at the same time")
 
-    query = session.query(Job).filter(Job.state == State.RUNNING).order_by(Job.latest_heartbeat.desc())
+    query = select(Job).where(Job.state == State.RUNNING).order_by(Job.latest_heartbeat.desc())
     if args.job_type:
-        query = query.filter(Job.job_type == args.job_type)
+        query = query.where(Job.job_type == args.job_type)
     if args.hostname:
-        query = query.filter(Job.hostname == args.hostname)
+        query = query.where(Job.hostname == args.hostname)
     if args.local:
-        query = query.filter(Job.hostname == get_hostname())
+        query = query.where(Job.hostname == get_hostname())
     if args.limit > 0:
         query = query.limit(args.limit)
 
-    alive_jobs: list[Job] = [job for job in query.all() if job.is_alive()]
+    alive_jobs: list[Job] = [job for job in session.scalars(query) if job.is_alive()]
 
     count_alive_jobs = len(alive_jobs)
     if count_alive_jobs == 0:

airflow/cli/commands/rotate_fernet_key_command.py

@@ -17,6 +17,8 @@
 """Rotate Fernet key command."""
 from __future__ import annotations
 
+from sqlalchemy import select
+
 from airflow.models import Connection, Variable
 from airflow.utils import cli as cli_utils
 from airflow.utils.session import create_session
@@ -26,7 +28,8 @@
 def rotate_fernet_key(args):
     """Rotates all encrypted connection credentials and variables."""
     with create_session() as session:
-        for conn in session.query(Connection).filter(Connection.is_encrypted | Connection.is_extra_encrypted):
+        conns_query = select(Connection).where(Connection.is_encrypted | Connection.is_extra_encrypted)
+        for conn in session.scalars(conns_query):
             conn.rotate_fernet_key()
-        for var in session.query(Variable).filter(Variable.is_encrypted):
+        for var in session.scalars(select(Variable).where(Variable.is_encrypted)):
             var.rotate_fernet_key()

airflow/cli/commands/task_command.py

@@ -29,6 +29,7 @@
 
 import pendulum
 from pendulum.parsing.exceptions import ParserError
+from sqlalchemy import select
 from sqlalchemy.orm.exc import NoResultFound
 from sqlalchemy.orm.session import Session
 
@@ -111,11 +112,9 @@ def _get_dag_run(
         with suppress(ParserError, TypeError):
             execution_date = timezone.parse(exec_date_or_run_id)
         try:
-            dag_run = (
-                session.query(DagRun)
-                .filter(DagRun.dag_id == dag.dag_id, DagRun.execution_date == execution_date)
-                .one()
-            )
+            dag_run = session.scalars(
+                select(DagRun).where(DagRun.dag_id == dag.dag_id, DagRun.execution_date == execution_date)
+            ).one()
         except NoResultFound:
             if not create_if_necessary:
                 raise DagRunNotFound(
@@ -534,18 +533,14 @@ def _guess_debugger() -> _SupportedDebugger:
 @provide_session
 def task_states_for_dag_run(args, session: Session = NEW_SESSION) -> None:
     """Get the status of all task instances in a DagRun."""
-    dag_run = (
-        session.query(DagRun)
-        .filter(DagRun.run_id == args.execution_date_or_run_id, DagRun.dag_id == args.dag_id)
-        .one_or_none()
+    dag_run = session.scalar(
+        select(DagRun).where(DagRun.run_id == args.execution_date_or_run_id, DagRun.dag_id == args.dag_id)
     )
     if not dag_run:
         try:
             execution_date = timezone.parse(args.execution_date_or_run_id)
-            dag_run = (
-                session.query(DagRun)
-                .filter(DagRun.execution_date == execution_date, DagRun.dag_id == args.dag_id)
-                .one_or_none()
+            dag_run = session.scalar(
+                select(DagRun).where(DagRun.execution_date == execution_date, DagRun.dag_id == args.dag_id)
             )
         except (ParserError, TypeError) as err:
             raise AirflowException(f"Error parsing the supplied execution_date. Error: {str(err)}")

airflow/cli/commands/variable_command.py

@@ -22,6 +22,8 @@
 import os
 from json import JSONDecodeError
 
+from sqlalchemy import select
+
 from airflow.cli.simple_table import AirflowConsole
 from airflow.models import Variable
 from airflow.utils import cli as cli_utils
@@ -33,7 +35,7 @@
 def variables_list(args):
     """Displays all the variables."""
     with create_session() as session:
-        variables = session.query(Variable)
+        variables = session.scalars(select(Variable)).all()
     AirflowConsole().print_as(data=variables, output=args.output, mapper=lambda x: {"key": x.key})
 
 
@@ -107,7 +109,7 @@ def _variable_export_helper(filepath):
     """Helps export all the variables to the file."""
     var_dict = {}
     with create_session() as session:
-        qry = session.query(Variable).all()
+        qry = session.scalars(select(Variable))
 
         data = json.JSONDecoder()
         for var in qry:

airflow/models/renderedtifields.py

@@ -134,15 +134,13 @@ def get_templated_fields(cls, ti: TaskInstance, session: Session = NEW_SESSION)
         :param session: SqlAlchemy Session
         :return: Rendered Templated TI field
         """
-        result = (
-            session.query(cls.rendered_fields)
-            .filter(
+        result = session.scalar(
+            select(cls).where(
                 cls.dag_id == ti.dag_id,
                 cls.task_id == ti.task_id,
                 cls.run_id == ti.run_id,
                 cls.map_index == ti.map_index,
             )
-            .one_or_none()
         )
 
         if result:
@@ -162,15 +160,13 @@ def get_k8s_pod_yaml(cls, ti: TaskInstance, session: Session = NEW_SESSION) -> d
         :param session: SqlAlchemy Session
         :return: Kubernetes Pod Yaml
         """
-        result = (
-            session.query(cls.k8s_pod_yaml)
-            .filter(
+        result = session.scalar(
+            select(cls).where(
                 cls.dag_id == ti.dag_id,
                 cls.task_id == ti.task_id,
                 cls.run_id == ti.run_id,
                 cls.map_index == ti.map_index,
             )
-            .one_or_none()
         )
         return result.k8s_pod_yaml if result else None
 
@@ -243,7 +239,8 @@ def _do_delete_old_records(
                 cls.task_id == task_id,
                 tuple_not_in_condition(
                     (cls.dag_id, cls.task_id, cls.run_id),
-                    session.query(ti_clause.c.dag_id, ti_clause.c.task_id, ti_clause.c.run_id),
+                    select(ti_clause.c.dag_id, ti_clause.c.task_id, ti_clause.c.run_id),
+                    session=session,
                 ),
             )
             .execution_options(synchronize_session=False)

airflow/utils/sqlalchemy.py

@@ -22,14 +22,14 @@
 import datetime
 import json
 import logging
-from typing import TYPE_CHECKING, Any, Generator, Iterable
+from typing import TYPE_CHECKING, Any, Generator, Iterable, overload
 
 import pendulum
 from dateutil import relativedelta
 from sqlalchemy import TIMESTAMP, PickleType, and_, event, false, nullsfirst, or_, true, tuple_
 from sqlalchemy.dialects import mssql, mysql
 from sqlalchemy.exc import OperationalError
-from sqlalchemy.sql import ColumnElement
+from sqlalchemy.sql import ColumnElement, Select
 from sqlalchemy.sql.expression import ColumnOperators
 from sqlalchemy.types import JSON, Text, TypeDecorator, TypeEngine, UnicodeText
 
@@ -515,11 +515,31 @@ def is_lock_not_available_error(error: OperationalError):
     return False
 
 
+@overload
 def tuple_in_condition(
     columns: tuple[ColumnElement, ...],
     collection: Iterable[Any],
 ) -> ColumnOperators:
-    """Generates a tuple-in-collection operator to use in ``.filter()``.
+    ...
+
+
+@overload
+def tuple_in_condition(
+    columns: tuple[ColumnElement, ...],
+    collection: Select,
+    *,
+    session: Session,
+) -> ColumnOperators:
+    ...
+
+
+def tuple_in_condition(
+    columns: tuple[ColumnElement, ...],
+    collection: Iterable[Any] | Select,
+    *,
+    session: Session | None = None,
+) -> ColumnOperators:
+    """Generates a tuple-in-collection operator to use in ``.where()``.
 
     For most SQL backends, this generates a simple ``([col, ...]) IN [condition]``
     clause. This however does not work with MSSQL, where we need to expand to
@@ -529,25 +549,57 @@ def tuple_in_condition(
     """
     if settings.engine.dialect.name != "mssql":
         return tuple_(*columns).in_(collection)
-    clauses = [and_(*(c == v for c, v in zip(columns, values))) for values in collection]
+    if not isinstance(collection, Select):
+        rows = collection
+    elif session is None:
+        raise TypeError("session is required when passing in a subquery")
+    else:
+        rows = session.execute(collection)
+    clauses = [and_(*(c == v for c, v in zip(columns, values))) for values in rows]
     if not clauses:
         return false()
     return or_(*clauses)
 
 
+@overload
 def tuple_not_in_condition(
     columns: tuple[ColumnElement, ...],
     collection: Iterable[Any],
 ) -> ColumnOperators:
-    """Generates a tuple-not-in-collection operator to use in ``.filter()``.
+    ...
+
+
+@overload
+def tuple_not_in_condition(
+    columns: tuple[ColumnElement, ...],
+    collection: Select,
+    *,
+    session: Session,
+) -> ColumnOperators:
+    ...
+
+
+def tuple_not_in_condition(
+    columns: tuple[ColumnElement, ...],
+    collection: Iterable[Any] | Select,
+    *,
+    session: Session | None = None,
+) -> ColumnOperators:
+    """Generates a tuple-not-in-collection operator to use in ``.where()``.
 
     This is similar to ``tuple_in_condition`` except generating ``NOT IN``.
 
     :meta private:
     """
     if settings.engine.dialect.name != "mssql":
         return tuple_(*columns).not_in(collection)
-    clauses = [or_(*(c != v for c, v in zip(columns, values))) for values in collection]
+    if not isinstance(collection, Select):
+        rows = collection
+    elif session is None:
+        raise TypeError("session is required when passing in a subquery")
+    else:
+        rows = session.execute(collection)
+    clauses = [or_(*(c != v for c, v in zip(columns, values))) for values in rows]
     if not clauses:
         return true()
     return and_(*clauses)

tests/cli/commands/test_task_command.py

@@ -500,15 +500,15 @@ def test_mapped_task_render_with_template(self, dag_maker):
         assert 'echo "2022-01-01"' in output
         assert 'echo "2022-01-08"' in output
 
-    @mock.patch("sqlalchemy.orm.session.Session.query")
+    @mock.patch("airflow.cli.commands.task_command.select")
+    @mock.patch("airflow.cli.commands.task_command.Session.scalars")
     @mock.patch("airflow.cli.commands.task_command.DagRun")
-    def test_task_render_with_custom_timetable(self, mock_dagrun, mock_query):
+    def test_task_render_with_custom_timetable(self, mock_dagrun, mock_scalars, mock_select):
         """
         when calling `tasks render` on dag with custom timetable, the DagRun object should be created with
          data_intervals.
         """
-        mock_query.side_effect = sqlalchemy.exc.NoResultFound
-
+        mock_scalars.side_effect = sqlalchemy.exc.NoResultFound
         task_command.task_render(
             self.parser.parse_args(["tasks", "render", "example_workday_timetable", "run_this", "2022-01-01"])
         )