From 65c7c30c5b560caa0fabb268be7a5aa70dee84b0 Mon Sep 17 00:00:00 2001
From: Phani Kumar <phani.kumar@astronomer.io>
Date: Fri, 21 Jul 2023 19:01:57 +0530
Subject: [PATCH 01/10] SqlAlchemy 2.0 changes

---
 airflow/api/common/experimental/pool.py       |  2 +-
 airflow/auth/managers/fab/models/__init__.py  | 14 ++--
 airflow/datasets/manager.py                   |  4 +-
 airflow/models/variable.py                    |  5 +-
 airflow/operators/subdag.py                   | 13 ++--
 .../celery/executors/celery_executor_utils.py |  3 +-
 airflow/ti_deps/deps/trigger_rule_dep.py      |  2 +-
 airflow/utils/scheduler_health.py             |  8 ++-
 airflow/www/security.py                       | 68 +++++++++----------
 tests/datasets/test_manager.py                |  2 +-
 10 files changed, 62 insertions(+), 59 deletions(-)

diff --git a/airflow/api/common/experimental/pool.py b/airflow/api/common/experimental/pool.py
index 34e35cd43556f..1134bda989014 100644
--- a/airflow/api/common/experimental/pool.py
+++ b/airflow/api/common/experimental/pool.py
@@ -45,7 +45,7 @@ def get_pool(name, session: Session = NEW_SESSION):
 @provide_session
 def get_pools(session: Session = NEW_SESSION):
     """Get all pools."""
-    return session.query(Pool).all()
+    return session.scalars(select(Pool)).all()
 
 
 @deprecated(reason="Use Pool.create_pool() instead", version="2.2.4")
diff --git a/airflow/auth/managers/fab/models/__init__.py b/airflow/auth/managers/fab/models/__init__.py
index 0bc26adb7eb4a..28de7840ec7ad 100644
--- a/airflow/auth/managers/fab/models/__init__.py
+++ b/airflow/auth/managers/fab/models/__init__.py
@@ -37,6 +37,7 @@
     UniqueConstraint,
     event,
     func,
+    select,
 )
 from sqlalchemy.orm import backref, declared_attr, relationship
 
@@ -210,12 +211,13 @@ def perms(self):
             if current_app:
                 sm = current_app.appbuilder.sm
                 self._perms: set[tuple[str, str]] = set(
-                    sm.get_session.query(sm.action_model.name, sm.resource_model.name)
-                    .join(sm.permission_model.action)
-                    .join(sm.permission_model.resource)
-                    .join(sm.permission_model.role)
-                    .filter(sm.role_model.user.contains(self))
-                    .all()
+                    sm.get_session.execute(
+                        select(sm.action_model.name, sm.resource_model.name)
+                        .join(sm.permission_model.action)
+                        .join(sm.permission_model.resource)
+                        .join(sm.permission_model.role)
+                        .where(sm.role_model.user.contains(self))
+                    )
                 )
             else:
                 self._perms = {
diff --git a/airflow/datasets/manager.py b/airflow/datasets/manager.py
index 2a1d55bba9156..c7e062803ab97 100644
--- a/airflow/datasets/manager.py
+++ b/airflow/datasets/manager.py
@@ -19,7 +19,7 @@
 
 from typing import TYPE_CHECKING
 
-from sqlalchemy import exc
+from sqlalchemy import exc, select
 from sqlalchemy.orm.session import Session
 
 from airflow.configuration import conf
@@ -52,7 +52,7 @@ def register_dataset_change(
         For local datasets, look them up, record the dataset event, queue dagruns, and broadcast
         the dataset event
         """
-        dataset_model = session.query(DatasetModel).filter(DatasetModel.uri == dataset.uri).one_or_none()
+        dataset_model = session.scalar(select(DatasetModel).where(DatasetModel.uri == dataset.uri))
         if not dataset_model:
             self.log.warning("DatasetModel %s not found", dataset)
             return
diff --git a/airflow/models/variable.py b/airflow/models/variable.py
index 5ca7774cbf064..07cc9ffb6c88f 100644
--- a/airflow/models/variable.py
+++ b/airflow/models/variable.py
@@ -21,7 +21,7 @@
 import logging
 from typing import Any
 
-from sqlalchemy import Boolean, Column, Integer, String, Text, delete
+from sqlalchemy import Boolean, Column, Integer, String, Text, delete, select
 from sqlalchemy.dialects.mysql import MEDIUMTEXT
 from sqlalchemy.orm import Session, declared_attr, reconstructor, synonym
 
@@ -200,8 +200,7 @@ def update(
 
         if Variable.get_variable_from_secrets(key=key) is None:
             raise KeyError(f"Variable {key} does not exist")
-
-        obj = session.query(Variable).filter(Variable.key == key).first()
+        obj = session.scalar(select(Variable).where(Variable.key == key))
         if obj is None:
             raise AttributeError(f"Variable {key} does not exist in the Database and cannot be updated.")
 
diff --git a/airflow/operators/subdag.py b/airflow/operators/subdag.py
index 680497217d51a..345c36e72b10c 100644
--- a/airflow/operators/subdag.py
+++ b/airflow/operators/subdag.py
@@ -26,6 +26,7 @@
 from datetime import datetime
 from enum import Enum
 
+from sqlalchemy import select
 from sqlalchemy.orm.session import Session
 
 from airflow.api.common.experimental.get_task_instance import get_task_instance
@@ -112,7 +113,7 @@ def _validate_pool(self, session):
             conflicts = [t for t in self.subdag.tasks if t.pool == self.pool]
             if conflicts:
                 # only query for pool conflicts if one may exist
-                pool = session.query(Pool).filter(Pool.slots == 1).filter(Pool.pool == self.pool).first()
+                pool = session.scalar(select(Pool).where(Pool.slots == 1, Pool.pool == self.pool))
                 if pool and any(t.pool == self.pool for t in self.subdag.tasks):
                     raise AirflowException(
                         f"SubDagOperator {self.task_id} and subdag task{'s' if len(conflicts) > 1 else ''} "
@@ -139,11 +140,11 @@ def _reset_dag_run_and_task_instances(self, dag_run: DagRun, execution_date: dat
         with create_session() as session:
             dag_run.state = DagRunState.RUNNING
             session.merge(dag_run)
-            failed_task_instances = (
-                session.query(TaskInstance)
-                .filter(TaskInstance.dag_id == self.subdag.dag_id)
-                .filter(TaskInstance.execution_date == execution_date)
-                .filter(TaskInstance.state.in_((TaskInstanceState.FAILED, TaskInstanceState.UPSTREAM_FAILED)))
+            failed_task_instances = session.scalars(
+                select(TaskInstance)
+                .where(TaskInstance.dag_id == self.subdag.dag_id)
+                .where(TaskInstance.execution_date == execution_date)
+                .where(TaskInstance.state.in_((TaskInstanceState.FAILED, TaskInstanceState.UPSTREAM_FAILED)))
             )
 
             for task_instance in failed_task_instances:
diff --git a/airflow/providers/celery/executors/celery_executor_utils.py b/airflow/providers/celery/executors/celery_executor_utils.py
index 2e739f239a434..5cd8ea7eb1768 100644
--- a/airflow/providers/celery/executors/celery_executor_utils.py
+++ b/airflow/providers/celery/executors/celery_executor_utils.py
@@ -37,6 +37,7 @@
 from celery.result import AsyncResult
 from celery.signals import import_modules as celery_import_modules
 from setproctitle import setproctitle
+from sqlalchemy import select
 
 import airflow.settings as settings
 from airflow.configuration import conf
@@ -268,7 +269,7 @@ def _query_task_cls_from_db_backend(self, task_ids, **kwargs):
         session = app.backend.ResultSession()
         task_cls = getattr(app.backend, "task_cls", TaskDb)
         with session_cleanup(session):
-            return session.query(task_cls).filter(task_cls.task_id.in_(task_ids)).all()
+            return session.scalars(select(task_cls).where(task_cls.task_id.in_(task_ids))).all()
 
     def _get_many_from_db_backend(self, async_tasks) -> Mapping[str, EventBufferValueType]:
         task_ids = self._tasks_list_to_task_ids(async_tasks)
diff --git a/airflow/ti_deps/deps/trigger_rule_dep.py b/airflow/ti_deps/deps/trigger_rule_dep.py
index c7e2982fffdd7..d9786f2c3b147 100644
--- a/airflow/ti_deps/deps/trigger_rule_dep.py
+++ b/airflow/ti_deps/deps/trigger_rule_dep.py
@@ -22,7 +22,7 @@
 import functools
 from typing import TYPE_CHECKING, Iterator, NamedTuple
 
-from sqlalchemy import and_, func, or_, select
+from sqlalchemy import and_, case, func, or_, select, true
 
 from airflow.models.taskinstance import PAST_DEPENDS_MET
 from airflow.ti_deps.dep_context import DepContext
diff --git a/airflow/utils/scheduler_health.py b/airflow/utils/scheduler_health.py
index 6108aee018a00..0a068d7f7293a 100644
--- a/airflow/utils/scheduler_health.py
+++ b/airflow/utils/scheduler_health.py
@@ -19,6 +19,8 @@
 import logging
 from http.server import BaseHTTPRequestHandler, HTTPServer
 
+from sqlalchemy import select
+
 from airflow.configuration import conf
 from airflow.jobs.job import Job
 from airflow.jobs.scheduler_job_runner import SchedulerJobRunner
@@ -35,12 +37,12 @@ def do_GET(self):
         if self.path == "/health":
             try:
                 with create_session() as session:
-                    scheduler_job = (
-                        session.query(Job)
+                    scheduler_job = session.scalar(
+                        select(Job)
                         .filter_by(job_type=SchedulerJobRunner.job_type)
                         .filter_by(hostname=get_hostname())
                         .order_by(Job.latest_heartbeat.desc())
-                        .first()
+                        .limit(1)
                     )
                 if scheduler_job and scheduler_job.is_alive():
                     self.send_response(200)
diff --git a/airflow/www/security.py b/airflow/www/security.py
index a988dccd2e32b..405aaea1c54c4 100644
--- a/airflow/www/security.py
+++ b/airflow/www/security.py
@@ -20,7 +20,7 @@
 from typing import TYPE_CHECKING, Any, Collection, Container, Iterable, Sequence
 
 from flask import g
-from sqlalchemy import or_
+from sqlalchemy import or_, select
 from sqlalchemy.orm import Session, joinedload
 
 from airflow.auth.managers.fab.models import Permission, Resource, Role, User
@@ -243,11 +243,9 @@ def __init__(self, appbuilder) -> None:
 
     def _get_root_dag_id(self, dag_id: str) -> str:
         if "." in dag_id:
-            dm = (
-                self.appbuilder.get_session.query(DagModel.dag_id, DagModel.root_dag_id)
-                .filter(DagModel.dag_id == dag_id)
-                .first()
-            )
+            dm = self.appbuilder.get_session.execute(
+                select(DagModel.dag_id, DagModel.root_dag_id).where(DagModel.dag_id == dag_id)
+            ).one()
             return dm.root_dag_id or dm.dag_id
         return dag_id
 
@@ -331,7 +329,7 @@ def get_accessible_dags(
             stacklevel=3,
         )
         dag_ids = self.get_accessible_dag_ids(user, user_actions, session)
-        return session.query(DagModel).filter(DagModel.dag_id.in_(dag_ids))
+        return session.scalars(select(DagModel).where(DagModel.dag_id.in_(dag_ids)))
 
     def get_readable_dag_ids(self, user) -> set[str]:
         """Gets the DAG IDs readable by authenticated user."""
@@ -358,16 +356,15 @@ def get_accessible_dag_ids(
             if (permissions.ACTION_CAN_EDIT in user_actions and self.can_edit_all_dags(user)) or (
                 permissions.ACTION_CAN_READ in user_actions and self.can_read_all_dags(user)
             ):
-                return {dag.dag_id for dag in session.query(DagModel.dag_id)}
-            user_query = (
-                session.query(User)
+                return {dag.dag_id for dag in session.execute(select(DagModel.dag_id))}
+            user_query = session.scalar(
+                select(User)
                 .options(
                     joinedload(User.roles)
                     .subqueryload(Role.permissions)
                     .options(joinedload(Permission.action), joinedload(Permission.resource))
                 )
-                .filter(User.id == user.id)
-                .first()
+                .where(User.id == user.id)
             )
             roles = user_query.roles
 
@@ -380,13 +377,16 @@ def get_accessible_dag_ids(
 
                 resource = permission.resource.name
                 if resource == permissions.RESOURCE_DAG:
-                    return {dag.dag_id for dag in session.query(DagModel.dag_id)}
+                    return {dag.dag_id for dag in session.execute(select(DagModel.dag_id))}
 
                 if resource.startswith(permissions.RESOURCE_DAG_PREFIX):
                     resources.add(resource[len(permissions.RESOURCE_DAG_PREFIX) :])
                 else:
                     resources.add(resource)
-        return {dag.dag_id for dag in session.query(DagModel.dag_id).filter(DagModel.dag_id.in_(resources))}
+        return {
+            dag.dag_id
+            for dag in session.execute(select(DagModel.dag_id).where(DagModel.dag_id.in_(resources)))
+        }
 
     def can_access_some_dags(self, action: str, dag_id: str | None = None) -> bool:
         """Checks if user has read or write access to some dags."""
@@ -524,10 +524,8 @@ def _merge_perm(self, action_name: str, resource_name: str) -> None:
         resource = self.get_resource(resource_name)
         perm = None
         if action and resource:
-            perm = (
-                self.appbuilder.get_session.query(self.permission_model)
-                .filter_by(action=action, resource=resource)
-                .first()
+            perm = self.appbuilder.get_session.scalar(
+                select(self.permission_model).filter_by(action=action, resource=resource)
             )
         if not perm and action_name and resource_name:
             self.create_permission(action_name, resource_name)
@@ -548,11 +546,11 @@ def add_homepage_access_to_custom_roles(self) -> None:
     def get_all_permissions(self) -> set[tuple[str, str]]:
         """Returns all permissions as a set of tuples with the action and resource names."""
         return set(
-            self.appbuilder.get_session.query(self.permission_model)
-            .join(self.permission_model.action)
-            .join(self.permission_model.resource)
-            .with_entities(self.action_model.name, self.resource_model.name)
-            .all()
+            self.appbuilder.get_session.execute(
+                select(self.action_model.name, self.resource_model.name)
+                .join(self.permission_model.action)
+                .join(self.permission_model.resource)
+            )
         )
 
     def _get_all_non_dag_permissions(self) -> dict[tuple[str, str], Permission]:
@@ -565,12 +563,12 @@ def _get_all_non_dag_permissions(self) -> dict[tuple[str, str], Permission]:
         return {
             (action_name, resource_name): viewmodel
             for action_name, resource_name, viewmodel in (
-                self.appbuilder.get_session.query(self.permission_model)
-                .join(self.permission_model.action)
-                .join(self.permission_model.resource)
-                .filter(~self.resource_model.name.like(f"{permissions.RESOURCE_DAG_PREFIX}%"))
-                .with_entities(self.action_model.name, self.resource_model.name, self.permission_model)
-                .all()
+                self.appbuilder.get_session.execute(
+                    select(self.action_model.name, self.resource_model.name, self.permission_model)
+                    .join(self.permission_model.action)
+                    .join(self.permission_model.resource)
+                    .where(~self.resource_model.name.like(f"{permissions.RESOURCE_DAG_PREFIX}%"))
+                )
             )
         }
 
@@ -578,9 +576,9 @@ def _get_all_roles_with_permissions(self) -> dict[str, Role]:
         """Returns a dict with a key of role name and value of role with early loaded permissions."""
         return {
             r.name: r
-            for r in self.appbuilder.get_session.query(self.role_model).options(
-                joinedload(self.role_model.permissions)
-            )
+            for r in self.appbuilder.get_session.scalars(
+                select(self.role_model).options(joinedload(self.role_model.permissions))
+            ).unique()
         }
 
     def create_dag_specific_permissions(self) -> None:
@@ -621,12 +619,12 @@ def update_admin_permission(self) -> None:
         :return: None.
         """
         session = self.appbuilder.get_session
-        dag_resources = session.query(Resource).filter(
-            Resource.name.like(f"{permissions.RESOURCE_DAG_PREFIX}%")
+        dag_resources = session.scalars(
+            select(Resource).where(Resource.name.like(f"{permissions.RESOURCE_DAG_PREFIX}%"))
         )
         resource_ids = [resource.id for resource in dag_resources]
 
-        perms = session.query(Permission).filter(~Permission.resource_id.in_(resource_ids))
+        perms = session.scalars(select(Permission).where(~Permission.resource_id.in_(resource_ids)))
         perms = [p for p in perms if p.action and p.resource]
 
         admin = self.find_role("Admin")
diff --git a/tests/datasets/test_manager.py b/tests/datasets/test_manager.py
index e462d0093e06c..10d54744b9f0e 100644
--- a/tests/datasets/test_manager.py
+++ b/tests/datasets/test_manager.py
@@ -54,7 +54,7 @@ def test_register_dataset_change_dataset_doesnt_exist(self, mock_task_instance):
 
         mock_session = mock.Mock()
         # Gotta mock up the query results
-        mock_session.query.return_value.filter.return_value.one_or_none.return_value = None
+        mock_session.scalar.return_value = None
 
         dsem.register_dataset_change(task_instance=mock_task_instance, dataset=dataset, session=mock_session)
 

From c015de26202e83ac3b13999cdbd6f3f4046ae44e Mon Sep 17 00:00:00 2001
From: Phani Kumar <phani.kumar@astronomer.io>
Date: Fri, 4 Aug 2023 15:27:29 +0530
Subject: [PATCH 02/10] Fix test failures

---
 airflow/ti_deps/deps/trigger_rule_dep.py            | 2 +-
 tests/integration/executors/test_celery_executor.py | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/airflow/ti_deps/deps/trigger_rule_dep.py b/airflow/ti_deps/deps/trigger_rule_dep.py
index d9786f2c3b147..c7e2982fffdd7 100644
--- a/airflow/ti_deps/deps/trigger_rule_dep.py
+++ b/airflow/ti_deps/deps/trigger_rule_dep.py
@@ -22,7 +22,7 @@
 import functools
 from typing import TYPE_CHECKING, Iterator, NamedTuple
 
-from sqlalchemy import and_, case, func, or_, select, true
+from sqlalchemy import and_, func, or_, select
 
 from airflow.models.taskinstance import PAST_DEPENDS_MET
 from airflow.ti_deps.dep_context import DepContext
diff --git a/tests/integration/executors/test_celery_executor.py b/tests/integration/executors/test_celery_executor.py
index 9e8d365c5ce80..26c2af03db8c6 100644
--- a/tests/integration/executors/test_celery_executor.py
+++ b/tests/integration/executors/test_celery_executor.py
@@ -311,7 +311,7 @@ def test_should_support_db_backend(self, mock_session, caplog):
             ):
                 caplog.clear()
                 mock_session = mock_backend.ResultSession.return_value
-                mock_session.query.return_value.filter.return_value.all.return_value = [
+                mock_session.scalars.return_value.all.return_value = [
                     mock.MagicMock(**{"to_dict.return_value": {"status": "SUCCESS", "task_id": "123"}})
                 ]
 
@@ -340,7 +340,7 @@ def test_should_retry_db_backend(self, mock_session, caplog):
             ):
                 caplog.clear()
                 mock_session = mock_backend.ResultSession.return_value
-                mock_retry_db_result = mock_session.query.return_value.filter.return_value.all
+                mock_retry_db_result = mock_session.scalars.return_value.all
                 mock_retry_db_result.return_value = [
                     mock.MagicMock(**{"to_dict.return_value": {"status": "SUCCESS", "task_id": "123"}})
                 ]

From 43c49f7f512104446c8a4276b414a464754a205a Mon Sep 17 00:00:00 2001
From: Phani Kumar <phani.kumar@astronomer.io>
Date: Mon, 7 Aug 2023 17:33:40 +0530
Subject: [PATCH 03/10] Change skipmixin

---
 airflow/models/skipmixin.py | 49 ++++++++++++++++++++-----------------
 1 file changed, 27 insertions(+), 22 deletions(-)

diff --git a/airflow/models/skipmixin.py b/airflow/models/skipmixin.py
index c8feb58d91d45..7cf393387d0c7 100644
--- a/airflow/models/skipmixin.py
+++ b/airflow/models/skipmixin.py
@@ -20,6 +20,8 @@
 import warnings
 from typing import TYPE_CHECKING, Iterable, Sequence
 
+from sqlalchemy import select, update
+
 from airflow.exceptions import AirflowException, RemovedInAirflow3Warning
 from airflow.models.dagrun import DagRun
 from airflow.models.taskinstance import TaskInstance
@@ -67,24 +69,29 @@ def _set_state_to_skipped(
         """Set state of task instances to skipped from the same dag run."""
         if tasks:
             now = timezone.utcnow()
-            TI = TaskInstance
-            query = session.query(TI).filter(
-                TI.dag_id == dag_run.dag_id,
-                TI.run_id == dag_run.run_id,
-            )
+
             if isinstance(tasks[0], tuple):
-                query = query.filter(tuple_in_condition((TI.task_id, TI.map_index), tasks))
+                session.execute(
+                    update(TaskInstance)
+                    .where(
+                        TaskInstance.dag_id == dag_run.dag_id,
+                        TaskInstance.run_id == dag_run.run_id,
+                        tuple_in_condition((TaskInstance.task_id, TaskInstance.map_index), tasks),
+                    )
+                    .values(state=TaskInstanceState.SKIPPED, start_date=now, end_date=now)
+                    .execution_options(synchronize_session=False)
+                )
             else:
-                query = query.filter(TI.task_id.in_(tasks))
-
-            query.update(
-                {
-                    TaskInstance.state: TaskInstanceState.SKIPPED,
-                    TaskInstance.start_date: now,
-                    TaskInstance.end_date: now,
-                },
-                synchronize_session=False,
-            )
+                session.execute(
+                    update(TaskInstance)
+                    .where(
+                        TaskInstance.dag_id == dag_run.dag_id,
+                        TaskInstance.run_id == dag_run.run_id,
+                        TaskInstance.task_id.in_(tasks),
+                    )
+                    .values(state=TaskInstanceState.SKIPPED, start_date=now, end_date=now)
+                    .execution_options(synchronize_session=False)
+                )
 
     @provide_session
     def skip(
@@ -121,14 +128,12 @@ def skip(
                 stacklevel=2,
             )
 
-            dag_run = (
-                session.query(DagRun)
-                .filter(
-                    DagRun.dag_id == task_list[0].dag_id,
-                    DagRun.execution_date == execution_date,
+            dag_run = session.scalar(
+                select(DagRun).where(
+                    DagRun.dag_id == task_list[0].dag_id, DagRun.execution_date == execution_date
                 )
-                .one()
             )
+
         elif execution_date and dag_run and execution_date != dag_run.execution_date:
             raise ValueError(
                 "execution_date has a different value to  dag_run.execution_date -- please only pass dag_run"

From f6499cf63a0fb5f8c11a1876760b4f1c3a9d9ea4 Mon Sep 17 00:00:00 2001
From: Phani Kumar <phani.kumar@astronomer.io>
Date: Tue, 8 Aug 2023 13:58:50 +0530
Subject: [PATCH 04/10] Changes to dag.py

---
 airflow/models/dag.py       | 3 +--
 airflow/utils/sqlalchemy.py | 2 +-
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/airflow/models/dag.py b/airflow/models/dag.py
index 75fee04145de7..80b29f8e4ce55 100644
--- a/airflow/models/dag.py
+++ b/airflow/models/dag.py
@@ -2029,8 +2029,7 @@ def set_task_group_state(
             raise ValueError("TaskGroup {group_id} could not be found")
         tasks_to_set_state = [task for task in task_group.iter_tasks() if isinstance(task, BaseOperator)]
         task_ids = [task.task_id for task in task_group.iter_tasks()]
-
-        dag_runs_query = session.query(DagRun.id).where(DagRun.dag_id == self.dag_id)
+        dag_runs_query = select(DagRun.id).where(DagRun.dag_id == self.dag_id)
         if start_date is None and end_date is None:
             dag_runs_query = dag_runs_query.where(DagRun.execution_date == start_date)
         else:
diff --git a/airflow/utils/sqlalchemy.py b/airflow/utils/sqlalchemy.py
index f97570b94c429..bb2277e4ed47b 100644
--- a/airflow/utils/sqlalchemy.py
+++ b/airflow/utils/sqlalchemy.py
@@ -436,7 +436,7 @@ def lock_rows(query: Query, session: Session) -> Generator[None, None, None]:
 
     :meta private:
     """
-    locked_rows = with_row_locks(query, session).all()
+    locked_rows = with_row_locks(query, session)
     yield
     del locked_rows
 

From ca211fa950ecc57c81917f89ce2c5b880e88bcc3 Mon Sep 17 00:00:00 2001
From: Phani Kumar <phani.kumar@astronomer.io>
Date: Thu, 10 Aug 2023 10:14:15 +0530
Subject: [PATCH 05/10] Changes to processor.py

---
 airflow/dag_processing/processor.py | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/airflow/dag_processing/processor.py b/airflow/dag_processing/processor.py
index 162fc5889c89d..498e8ccefa166 100644
--- a/airflow/dag_processing/processor.py
+++ b/airflow/dag_processing/processor.py
@@ -30,7 +30,7 @@
 from typing import TYPE_CHECKING, Iterable, Iterator
 
 from setproctitle import setproctitle
-from sqlalchemy import delete, exc, func, or_
+from sqlalchemy import delete, exc, func, or_, select
 from sqlalchemy.orm.session import Session
 
 from airflow import settings
@@ -428,20 +428,20 @@ def manage_slas(cls, dag_folder, dag_id: str, session: Session = NEW_SESSION) ->
         if not any(isinstance(ti.sla, timedelta) for ti in dag.tasks):
             cls.logger().info("Skipping SLA check for %s because no tasks in DAG have SLAs", dag)
             return
-
         qry = (
-            session.query(TI.task_id, func.max(DR.execution_date).label("max_ti"))
+            select(TI.task_id, func.max(DR.execution_date).label("max_ti"))
             .join(TI.dag_run)
-            .filter(TI.dag_id == dag.dag_id)
-            .filter(or_(TI.state == TaskInstanceState.SUCCESS, TI.state == TaskInstanceState.SKIPPED))
-            .filter(TI.task_id.in_(dag.task_ids))
+            .where(TI.dag_id == dag.dag_id)
+            .where(or_(TI.state == TaskInstanceState.SUCCESS, TI.state == TaskInstanceState.SKIPPED))
+            .where(TI.task_id.in_(dag.task_ids))
             .group_by(TI.task_id)
-            .subquery("sq")
         )
         # get recorded SlaMiss
         recorded_slas_query = set(
-            session.query(SlaMiss.dag_id, SlaMiss.task_id, SlaMiss.execution_date).filter(
-                SlaMiss.dag_id == dag.dag_id, SlaMiss.task_id.in_(dag.task_ids)
+            session.scalars(
+                select(SlaMiss.dag_id, SlaMiss.task_id, SlaMiss.execution_date).where(
+                    SlaMiss.dag_id == dag.dag_id, SlaMiss.task_id.in_(dag.task_ids)
+                )
             )
         )
 

From beb1dc2a16d6bb1b8d77c16265fc6bdc2af2c786 Mon Sep 17 00:00:00 2001
From: Phani Kumar <phani.kumar@astronomer.io>
Date: Fri, 11 Aug 2023 15:01:53 +0530
Subject: [PATCH 06/10] Fixup processor.py

---
 airflow/dag_processing/processor.py | 33 +++++++++++------------------
 1 file changed, 12 insertions(+), 21 deletions(-)

diff --git a/airflow/dag_processing/processor.py b/airflow/dag_processing/processor.py
index 498e8ccefa166..fad9849511cfa 100644
--- a/airflow/dag_processing/processor.py
+++ b/airflow/dag_processing/processor.py
@@ -435,25 +435,21 @@ def manage_slas(cls, dag_folder, dag_id: str, session: Session = NEW_SESSION) ->
             .where(or_(TI.state == TaskInstanceState.SUCCESS, TI.state == TaskInstanceState.SKIPPED))
             .where(TI.task_id.in_(dag.task_ids))
             .group_by(TI.task_id)
+            .subquery("sq")
         )
         # get recorded SlaMiss
         recorded_slas_query = set(
-            session.scalars(
+            session.execute(
                 select(SlaMiss.dag_id, SlaMiss.task_id, SlaMiss.execution_date).where(
                     SlaMiss.dag_id == dag.dag_id, SlaMiss.task_id.in_(dag.task_ids)
                 )
             )
         )
-
-        max_tis: Iterator[TI] = (
-            session.query(TI)
+        max_tis: Iterator[TI] = session.scalars(
+            select(TI)
             .join(TI.dag_run)
-            .filter(
-                TI.dag_id == dag.dag_id,
-                TI.task_id == qry.c.task_id,
-                DR.execution_date == qry.c.max_ti,
-            )
-        )
+            .where(TI.dag_id == dag.dag_id, TI.task_id == qry.c.task_id, DR.execution_date == qry.c.max_ti)
+        ).all()
 
         ts = timezone.utcnow()
 
@@ -490,23 +486,18 @@ def manage_slas(cls, dag_folder, dag_id: str, session: Session = NEW_SESSION) ->
             if sla_misses:
                 session.add_all(sla_misses)
         session.commit()
-
-        slas: list[SlaMiss] = (
-            session.query(SlaMiss)
-            .filter(SlaMiss.notification_sent == False, SlaMiss.dag_id == dag.dag_id)  # noqa
-            .all()
-        )
+        slas: list[SlaMiss] = session.scalars(
+            select(SlaMiss).where(SlaMiss.notification_sent is False, SlaMiss.dag_id == dag.dag_id)
+        ).all()
         if slas:
             sla_dates: list[datetime] = [sla.execution_date for sla in slas]
-            fetched_tis: list[TI] = (
-                session.query(TI)
-                .filter(
+            fetched_tis: list[TI] = session.scalars(
+                select(TI).where(
                     TI.dag_id == dag.dag_id,
                     TI.execution_date.in_(sla_dates),
                     TI.state != TaskInstanceState.SUCCESS,
                 )
-                .all()
-            )
+            ).all()
             blocking_tis: list[TI] = []
             for ti in fetched_tis:
                 if ti.task_id in dag.task_ids:

From 1f6534dfee7e4b3fd28f2d7a9ef8e06e01fb450a Mon Sep 17 00:00:00 2001
From: Phani Kumar <phani.kumar@astronomer.io>
Date: Mon, 14 Aug 2023 18:10:22 +0530
Subject: [PATCH 07/10] Fixup processor.py

---
 airflow/dag_processing/processor.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/airflow/dag_processing/processor.py b/airflow/dag_processing/processor.py
index fad9849511cfa..64858be94ed9e 100644
--- a/airflow/dag_processing/processor.py
+++ b/airflow/dag_processing/processor.py
@@ -449,7 +449,7 @@ def manage_slas(cls, dag_folder, dag_id: str, session: Session = NEW_SESSION) ->
             select(TI)
             .join(TI.dag_run)
             .where(TI.dag_id == dag.dag_id, TI.task_id == qry.c.task_id, DR.execution_date == qry.c.max_ti)
-        ).all()
+        )
 
         ts = timezone.utcnow()
 
@@ -487,7 +487,7 @@ def manage_slas(cls, dag_folder, dag_id: str, session: Session = NEW_SESSION) ->
                 session.add_all(sla_misses)
         session.commit()
         slas: list[SlaMiss] = session.scalars(
-            select(SlaMiss).where(SlaMiss.notification_sent is False, SlaMiss.dag_id == dag.dag_id)
+            select(SlaMiss).where(~SlaMiss.notification_sent, SlaMiss.dag_id == dag.dag_id)
         ).all()
         if slas:
             sla_dates: list[datetime] = [sla.execution_date for sla in slas]

From 0f518d2c3fdcc6054a55f477252e2fe342cb6da5 Mon Sep 17 00:00:00 2001
From: Phani Kumar <phani.kumar@astronomer.io>
Date: Wed, 16 Aug 2023 19:42:09 +0530
Subject: [PATCH 08/10] Change kubernetes executor

---
 .../executors/kubernetes_executor.py          | 25 +++++++++++--------
 1 file changed, 15 insertions(+), 10 deletions(-)

diff --git a/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py b/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py
index cf74c34d97405..8e95f7a6c9387 100644
--- a/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py
+++ b/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py
@@ -34,6 +34,7 @@
 from queue import Empty, Queue
 from typing import TYPE_CHECKING, Any, Sequence
 
+from sqlalchemy import select, update
 from sqlalchemy.orm import Session
 
 from airflow import AirflowException
@@ -215,12 +216,12 @@ def clear_not_launched_queued_tasks(self, session: Session = NEW_SESSION) -> Non
         from airflow.models.taskinstance import TaskInstance
 
         self.log.debug("Clearing tasks that have not been launched")
-        query = session.query(TaskInstance).filter(
+        query = select(TaskInstance).where(
             TaskInstance.state == TaskInstanceState.QUEUED, TaskInstance.queued_by_job_id == self.job_id
         )
         if self.kubernetes_queue:
-            query = query.filter(TaskInstance.queue == self.kubernetes_queue)
-        queued_tis: list[TaskInstance] = query.all()
+            query = query.where(TaskInstance.queue == self.kubernetes_queue)
+        queued_tis: list[TaskInstance] = session.scalars(query).all()
         self.log.info("Found %s queued task instances", len(queued_tis))
 
         # Go through the "last seen" dictionary and clean out old entries
@@ -262,12 +263,16 @@ def clear_not_launched_queued_tasks(self, session: Session = NEW_SESSION) -> Non
             if pod_list:
                 continue
             self.log.info("TaskInstance: %s found in queued state but was not launched, rescheduling", ti)
-            session.query(TaskInstance).filter(
-                TaskInstance.dag_id == ti.dag_id,
-                TaskInstance.task_id == ti.task_id,
-                TaskInstance.run_id == ti.run_id,
-                TaskInstance.map_index == ti.map_index,
-            ).update({TaskInstance.state: TaskInstanceState.SCHEDULED})
+            session.execute(
+                update(TaskInstance)
+                .where(
+                    TaskInstance.dag_id == ti.dag_id,
+                    TaskInstance.task_id == ti.task_id,
+                    TaskInstance.run_id == ti.run_id,
+                    TaskInstance.map_index == ti.map_index,
+                )
+                .values(state=TaskInstanceState.SCHEDULED)
+            )
 
     def start(self) -> None:
         """Starts the executor."""
@@ -457,7 +462,7 @@ def _change_state(
         if state is None:
             from airflow.models.taskinstance import TaskInstance
 
-            state = session.query(TaskInstance.state).filter(TaskInstance.filter_for_tis([key])).scalar()
+            state = session.scalar(select(TaskInstance.state).where(TaskInstance.filter_for_tis([key])))
             state = TaskInstanceState(state)
 
         self.event_buffer[key] = state, None

From 935da4d96bf38b0d94163978ee98bf0b63efec2a Mon Sep 17 00:00:00 2001
From: Phani Kumar <phani.kumar@astronomer.io>
Date: Sat, 19 Aug 2023 09:08:13 +0530
Subject: [PATCH 09/10] Apply review suggestions

---
 airflow/models/skipmixin.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/airflow/models/skipmixin.py b/airflow/models/skipmixin.py
index 7cf393387d0c7..0616d042b950c 100644
--- a/airflow/models/skipmixin.py
+++ b/airflow/models/skipmixin.py
@@ -128,11 +128,11 @@ def skip(
                 stacklevel=2,
             )
 
-            dag_run = session.scalar(
+            dag_run = session.scalars(
                 select(DagRun).where(
                     DagRun.dag_id == task_list[0].dag_id, DagRun.execution_date == execution_date
                 )
-            )
+            ).one()
 
         elif execution_date and dag_run and execution_date != dag_run.execution_date:
             raise ValueError(

From eacdeb4ed967313ba49721944c95a3721b5d7afb Mon Sep 17 00:00:00 2001
From: Phani Kumar <phani.kumar@astronomer.io>
Date: Mon, 21 Aug 2023 15:39:36 +0530
Subject: [PATCH 10/10] Apply review suggestions

---
 airflow/www/security.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/airflow/www/security.py b/airflow/www/security.py
index 405aaea1c54c4..93f8fd11fcbec 100644
--- a/airflow/www/security.py
+++ b/airflow/www/security.py
@@ -525,7 +525,7 @@ def _merge_perm(self, action_name: str, resource_name: str) -> None:
         perm = None
         if action and resource:
             perm = self.appbuilder.get_session.scalar(
-                select(self.permission_model).filter_by(action=action, resource=resource)
+                select(self.permission_model).filter_by(action=action, resource=resource).limit(1)
             )
         if not perm and action_name and resource_name:
             self.create_permission(action_name, resource_name)
