From 65c7c30c5b560caa0fabb268be7a5aa70dee84b0 Mon Sep 17 00:00:00 2001 From: Phani Kumar 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 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 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 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 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 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 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 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 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 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)