This is an automated email from the ASF dual-hosted git repository.
vincbeck pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 06ca8a9c248 Remove query obj from providers fab (#53953)
06ca8a9c248 is described below
commit 06ca8a9c248b69112b3a916f2458b1be02bbc07b
Author: Pratiksha <[email protected]>
AuthorDate: Sat Aug 30 00:16:58 2025 +0530
Remove query obj from providers fab (#53953)
* remove session.query from tests
* remove session.query from fab/src
* fix errors
* Removed unnecessary comments code and applied small fixes
* resolved conflict and removed session.query from it
* remove .query()
* fix static checks
* fix static checks
* fix import error
* include fab provider in pre-commit check
* remove query obj from test_user_endpoint
---
.pre-commit-config.yaml | 1 +
.../fab/auth_manager/security_manager/override.py | 89 ++++++++++++----------
.../test_role_and_permission_endpoint.py | 5 +-
.../api_endpoints/test_user_endpoint.py | 33 ++++----
.../fab/auth_manager/schemas/test_user_schema.py | 7 +-
.../tests/unit/fab/auth_manager/test_security.py | 13 ++--
providers/fab/tests/unit/fab/utils.py | 40 ++++------
.../fab/www/views/test_views_custom_user_views.py | 19 ++---
8 files changed, 109 insertions(+), 98 deletions(-)
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 6947bb96dd7..35c062e6e25 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -534,6 +534,7 @@ repos:
files: >
(?x)
^airflow-ctl.*\.py$|
+ ^providers/fab.*\.py$|
^task_sdk.*\.py$
pass_filenames: true
- id: update-supported-versions
diff --git
a/providers/fab/src/airflow/providers/fab/auth_manager/security_manager/override.py
b/providers/fab/src/airflow/providers/fab/auth_manager/security_manager/override.py
index d474335b833..2fa5c0c5b4b 100644
---
a/providers/fab/src/airflow/providers/fab/auth_manager/security_manager/override.py
+++
b/providers/fab/src/airflow/providers/fab/auth_manager/security_manager/override.py
@@ -63,7 +63,7 @@ from flask_jwt_extended import JWTManager
from flask_login import LoginManager
from itsdangerous import want_bytes
from markupsafe import Markup, escape
-from sqlalchemy import func, inspect, or_, select
+from sqlalchemy import delete, func, inspect, or_, select
from sqlalchemy.exc import MultipleResultsFound
from sqlalchemy.orm import joinedload
from werkzeug.security import check_password_hash, generate_password_hash
@@ -565,7 +565,7 @@ class
FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
interface = self.appbuilder.get_app.session_interface
session = interface.db.session
user_session_model = interface.sql_session_model
- num_sessions = session.query(user_session_model).count()
+ num_sessions =
session.scalars(select(func.count()).select_from(user_session_model)).one()
if num_sessions > MAX_NUM_DATABASE_USER_SESSIONS:
safe_username = escape(user.username)
self._cli_safe_flash(
@@ -580,7 +580,7 @@ class
FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
"warning",
)
else:
- for s in session.query(user_session_model):
+ for s in session.scalars(select(user_session_model)).all():
session_details =
interface.serializer.loads(want_bytes(s.data))
if session_details.get("_user_id") == user.id:
session.delete(s)
@@ -1209,12 +1209,14 @@ class
FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
"""FAB leaves faulty permissions that need to be cleaned up."""
self.log.debug("Cleaning faulty perms")
sesh = self.appbuilder.get_session
- perms = sesh.query(Permission).filter(
- or_(
- Permission.action == None, # noqa: E711
- Permission.resource == None, # noqa: E711
+ perms = sesh.scalars(
+ select(Permission).where(
+ or_(
+ Permission.action == None, # noqa: E711
+ Permission.resource == None, # noqa: E711
+ )
)
- )
+ ).all()
# Since FAB doesn't define ON DELETE CASCADE on these tables, we need
# to delete the _object_ so that SQLA knows to delete the many-to-many
# relationship object too. :(
@@ -1292,10 +1294,10 @@ class
FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
:param name: the role name
"""
- return
self.get_session.query(self.role_model).filter_by(name=name).one_or_none()
+ return
self.get_session.scalars(select(self.role_model).filter_by(name=name)).unique().one_or_none()
def get_all_roles(self):
- return self.get_session.query(self.role_model).all()
+ return self.get_session.scalars(select(self.role_model)).unique().all()
def delete_role(self, role_name: str) -> None:
"""
@@ -1304,10 +1306,10 @@ class
FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
:param role_name: the name of a role in the ab_role table
"""
session = self.get_session
- role = session.query(Role).filter(Role.name == role_name).first()
+ role = session.scalars(select(Role).where(Role.name ==
role_name)).first()
if role:
log.info("Deleting role '%s'", role_name)
- session.delete(role)
+ session.execute(delete(Role).where(Role.name == role_name))
session.commit()
else:
raise AirflowException(f"Role named '{role_name}' does not exist")
@@ -1338,7 +1340,11 @@ class
FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
return _roles
def get_public_role(self):
- return
self.get_session.query(self.role_model).filter_by(name=self.auth_role_public).one_or_none()
+ return (
+
self.get_session.scalars(select(self.role_model).filter_by(name=self.auth_role_public))
+ .unique()
+ .one_or_none()
+ )
"""
-----------
@@ -1395,7 +1401,7 @@ class
FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
def count_users(self):
"""Return the number of users in the database."""
- return self.get_session.query(func.count(self.user_model.id)).scalar()
+ return self.get_session.scalar(select(func.count(self.user_model.id)))
def add_register_user(self, username, first_name, last_name, email,
password="", hashed_password=""):
"""
@@ -1427,22 +1433,22 @@ class
FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
if username:
try:
if self.auth_username_ci:
- return (
- self.get_session.query(self.user_model)
- .filter(func.lower(self.user_model.username) ==
func.lower(username))
- .one_or_none()
+ return self.get_session.scalars(
+ select(self.user_model).where(
+ func.lower(self.user_model.username) ==
func.lower(username)
+ )
+ ).one_or_none()
+ return self.get_session.scalars(
+ select(self.user_model).where(
+ func.lower(self.user_model.username) ==
func.lower(username)
)
- return (
- self.get_session.query(self.user_model)
- .filter(func.lower(self.user_model.username) ==
func.lower(username))
- .one_or_none()
- )
+ ).one_or_none()
except MultipleResultsFound:
log.error("Multiple results found for user %s", username)
return None
elif email:
try:
- return
self.get_session.query(self.user_model).filter_by(email=email).one_or_none()
+ return
self.get_session.scalars(select(self.user_model).filter_by(email=email)).one_or_none()
except MultipleResultsFound:
log.error("Multiple results found for user with email %s",
email)
return None
@@ -1474,7 +1480,7 @@ class
FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
return False
def get_all_users(self):
- return self.get_session.query(self.user_model).all()
+ return self.get_session.scalars(select(self.user_model)).all()
def update_user_auth_stat(self, user, success=True):
"""
@@ -1514,7 +1520,7 @@ class
FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
:param name: name
"""
- return
self.get_session.query(self.action_model).filter_by(name=name).one_or_none()
+ return
self.get_session.scalars(select(self.action_model).filter_by(name=name)).one_or_none()
def create_action(self, name):
"""
@@ -1547,11 +1553,9 @@ class
FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
log.warning(const.LOGMSG_WAR_SEC_DEL_PERMISSION, name)
return False
try:
- perms = (
- self.get_session.query(self.permission_model)
- .filter(self.permission_model.action == action)
- .all()
- )
+ perms = self.get_session.scalars(
+
select(self.permission_model).where(self.permission_model.action == action)
+ ).all()
if perms:
log.warning(const.LOGMSG_WAR_SEC_DEL_PERM_PVM, action, perms)
return False
@@ -1575,7 +1579,7 @@ class
FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
:param name: Name of resource
"""
- return
self.get_session.query(self.resource_model).filter_by(name=name).one_or_none()
+ return
self.get_session.scalars(select(self.resource_model).filter_by(name=name)).one_or_none()
def create_resource(self, name) -> Resource | None:
"""
@@ -1617,10 +1621,13 @@ class
FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
resource = self.get_resource(resource_name)
if action and resource:
return (
- self.get_session.query(self.permission_model)
- .filter_by(action=action, resource=resource)
+ self.get_session.scalars(
+ select(self.permission_model).filter_by(action=action,
resource=resource)
+ )
+ .unique()
.one_or_none()
)
+
return None
def get_resource_permissions(self, resource: Resource) -> Permission:
@@ -1629,7 +1636,9 @@ class
FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
:param resource: Object representing a single resource.
"""
- return
self.get_session.query(self.permission_model).filter_by(resource_id=resource.id).all()
+ return self.get_session.scalars(
+ select(self.permission_model).filter_by(resource_id=resource.id)
+ ).all()
def create_permission(self, action_name, resource_name) -> Permission |
None:
"""
@@ -1676,9 +1685,9 @@ class
FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
perm = self.get_permission(action_name, resource_name)
if not perm:
return
- roles = (
-
self.get_session.query(self.role_model).filter(self.role_model.permissions.contains(perm)).first()
- )
+ roles = self.get_session.scalars(
+
select(self.role_model).where(self.role_model.permissions.contains(perm))
+ ).first()
if roles:
log.warning(const.LOGMSG_WAR_SEC_DEL_PERMVIEW, resource_name,
action_name, roles)
return
@@ -1687,7 +1696,9 @@ class
FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
self.get_session.delete(perm)
self.get_session.commit()
# if no more permission on permission view, delete permission
- if not
self.get_session.query(self.permission_model).filter_by(action=perm.action).all():
+ if not self.get_session.scalars(
+ select(self.permission_model).filter_by(action=perm.action)
+ ).all():
self.delete_action(perm.action.name)
log.info(const.LOGMSG_INF_SEC_DEL_PERMVIEW, action_name,
resource_name)
except Exception as e:
diff --git
a/providers/fab/tests/unit/fab/auth_manager/api_endpoints/test_role_and_permission_endpoint.py
b/providers/fab/tests/unit/fab/auth_manager/api_endpoints/test_role_and_permission_endpoint.py
index 8672456fb7a..76dd1923779 100644
---
a/providers/fab/tests/unit/fab/auth_manager/api_endpoints/test_role_and_permission_endpoint.py
+++
b/providers/fab/tests/unit/fab/auth_manager/api_endpoints/test_role_and_permission_endpoint.py
@@ -17,6 +17,7 @@
from __future__ import annotations
import pytest
+from sqlalchemy import select
from airflow.providers.fab.www.api_connexion.exceptions import
EXCEPTIONS_LINK_MAP
from airflow.providers.fab.www.security import permissions
@@ -73,7 +74,7 @@ class TestRoleEndpoint:
session = self.app.appbuilder.get_session
existing_roles = set(EXISTING_ROLES)
existing_roles.update(["Test", "TestNoPermissions"])
- roles =
session.query(Role).filter(~Role.name.in_(existing_roles)).all()
+ roles =
session.scalars(select(Role).where(~Role.name.in_(existing_roles))).unique().all()
for role in roles:
delete_role(self.app, role.name)
@@ -353,7 +354,7 @@ class TestDeleteRole(TestRoleEndpoint):
role = create_role(self.app, "mytestrole")
response = self.client.delete(f"/fab/v1/roles/{role.name}",
environ_overrides={"REMOTE_USER": "test"})
assert response.status_code == 204
- role_obj = session.query(Role).filter(Role.name == role.name).all()
+ role_obj = session.scalars(select(Role).where(Role.name ==
role.name)).all()
assert len(role_obj) == 0
def test_delete_should_respond_404(self):
diff --git
a/providers/fab/tests/unit/fab/auth_manager/api_endpoints/test_user_endpoint.py
b/providers/fab/tests/unit/fab/auth_manager/api_endpoints/test_user_endpoint.py
index 8ac12eed715..0b8b6eb2913 100644
---
a/providers/fab/tests/unit/fab/auth_manager/api_endpoints/test_user_endpoint.py
+++
b/providers/fab/tests/unit/fab/auth_manager/api_endpoints/test_user_endpoint.py
@@ -19,11 +19,10 @@ from __future__ import annotations
import unittest.mock
import pytest
-from sqlalchemy.sql.functions import count
+from sqlalchemy import delete, func, select
from airflow.providers.fab.www.api_connexion.exceptions import
EXCEPTIONS_LINK_MAP
from airflow.providers.fab.www.security import permissions
-from airflow.utils import timezone
from airflow.utils.session import create_session
from tests_common.test_utils.compat import ignore_provider_compatibility_error
@@ -35,6 +34,12 @@ from unit.fab.auth_manager.api_endpoints.api_connexion_utils
import (
delete_user,
)
+try:
+ from airflow.utils import timezone # type: ignore[attr-defined]
+except AttributeError:
+ from airflow.sdk import timezone
+
+
with ignore_provider_compatibility_error("2.9.0+", __file__):
from airflow.providers.fab.auth_manager.models import User
@@ -85,8 +90,7 @@ class TestUserEndpoint:
def teardown_method(self) -> None:
# Delete users that have our custom default time
- users = self.session.query(User).filter(User.changed_on ==
timezone.parse(DEFAULT_TIME))
- users.delete(synchronize_session=False)
+ self.session.execute(delete(User).where(User.changed_on ==
timezone.parse(DEFAULT_TIME)))
self.session.commit()
def _create_users(self, count, roles=None):
@@ -354,11 +358,11 @@ EXAMPLE_USER_EMAIL = "[email protected]"
def _delete_user(**filters):
with create_session() as session:
- user = session.query(User).filter_by(**filters).first()
+ user = session.scalars(select(User).filter_by(**filters)).first()
if user is None:
return
user.roles = []
- session.delete(user)
+ session.execute(delete(User).filter_by(**filters))
@pytest.fixture
@@ -676,10 +680,7 @@ class TestPatchUser(TestUserEndpoint):
assert "password" not in response.json
mock_generate_password_hash.assert_called_once_with("new-pass")
-
- password_in_db = (
- self.session.query(User.password).filter(User.username ==
autoclean_username).scalar()
- )
+ password_in_db =
self.session.scalar(select(User.password).where(User.username ==
autoclean_username))
assert password_in_db == "fake-hashed-pass"
@pytest.mark.usefixtures("autoclean_admin_user")
@@ -788,7 +789,9 @@ class TestDeleteUser(TestUserEndpoint):
environ_overrides={"REMOTE_USER": "test"},
)
assert response.status_code == 204, response.json # NO CONTENT.
- assert self.session.query(count(User.id)).filter(User.username ==
autoclean_username).scalar() == 0
+ assert (
+
self.session.scalar(select(func.count(User.id)).where(User.username ==
autoclean_username)) == 0
+ )
@pytest.mark.usefixtures("autoclean_admin_user")
def test_unauthenticated(self, autoclean_username):
@@ -796,7 +799,9 @@ class TestDeleteUser(TestUserEndpoint):
f"/fab/v1/users/{autoclean_username}",
)
assert response.status_code == 401, response.json
- assert self.session.query(count(User.id)).filter(User.username ==
autoclean_username).scalar() == 1
+ assert (
+
self.session.scalar(select(func.count(User.id)).where(User.username ==
autoclean_username)) == 1
+ )
@pytest.mark.usefixtures("autoclean_admin_user")
def test_forbidden(self, autoclean_username):
@@ -805,7 +810,9 @@ class TestDeleteUser(TestUserEndpoint):
environ_overrides={"REMOTE_USER": "test_no_permissions"},
)
assert response.status_code == 403, response.json
- assert self.session.query(count(User.id)).filter(User.username ==
autoclean_username).scalar() == 1
+ assert (
+
self.session.scalar(select(func.count(User.id)).where(User.username ==
autoclean_username)) == 1
+ )
def test_not_found(self, autoclean_username):
# This test does not populate autoclean_admin_user into the database.
diff --git
a/providers/fab/tests/unit/fab/auth_manager/schemas/test_user_schema.py
b/providers/fab/tests/unit/fab/auth_manager/schemas/test_user_schema.py
index 7b1c06e7017..91853b16eb3 100644
--- a/providers/fab/tests/unit/fab/auth_manager/schemas/test_user_schema.py
+++ b/providers/fab/tests/unit/fab/auth_manager/schemas/test_user_schema.py
@@ -17,6 +17,7 @@
from __future__ import annotations
import pytest
+from sqlalchemy import select
from airflow.utils import timezone
@@ -60,7 +61,7 @@ class TestUserBase:
self.session = self.app.appbuilder.get_session
def teardown_method(self):
- user = self.session.query(User).filter(User.email ==
TEST_EMAIL).first()
+ user = self.session.scalars(select(User).where(User.email ==
TEST_EMAIL)).first()
if user:
self.session.delete(user)
self.session.commit()
@@ -80,7 +81,7 @@ class TestUserCollectionItemSchema(TestUserBase):
self.session.add(user_model)
user_model.roles = [self.role]
self.session.commit()
- user = self.session.query(User).filter(User.email ==
TEST_EMAIL).first()
+ user = self.session.scalars(select(User).where(User.email ==
TEST_EMAIL)).first()
deserialized_user = user_collection_item_schema.dump(user)
# No user_id and password in dump
assert deserialized_user == {
@@ -111,7 +112,7 @@ class TestUserSchema(TestUserBase):
)
self.session.add(user_model)
self.session.commit()
- user = self.session.query(User).filter(User.email ==
TEST_EMAIL).first()
+ user = self.session.scalars(select(User).where(User.email ==
TEST_EMAIL)).first()
deserialized_user = user_schema.dump(user)
# No user_id and password in dump
assert deserialized_user == {
diff --git a/providers/fab/tests/unit/fab/auth_manager/test_security.py
b/providers/fab/tests/unit/fab/auth_manager/test_security.py
index 3c1c797123f..8448239a4ce 100644
--- a/providers/fab/tests/unit/fab/auth_manager/test_security.py
+++ b/providers/fab/tests/unit/fab/auth_manager/test_security.py
@@ -28,7 +28,7 @@ import pytest
import time_machine
from flask_appbuilder import SQLA, Model, expose, has_access
from flask_appbuilder.views import BaseView, ModelView
-from sqlalchemy import Column, Date, Float, Integer, String
+from sqlalchemy import Column, Date, Float, Integer, String, delete, func,
select
from airflow.exceptions import AirflowException
from airflow.models import DagModel
@@ -138,7 +138,7 @@ def _create_dag_bundle(bundle_name, session):
def _delete_dag_bundle(bundle_name, session):
- session.query(DagBundleModel).filter(DagBundleModel.name ==
bundle_name).delete()
+ session.execute(delete(DagBundleModel).where(DagBundleModel.name ==
bundle_name))
session.commit()
@@ -941,9 +941,10 @@ def test_no_additional_dag_permission_views_created(db,
security_manager):
ab_perm_role = assoc_permission_role
security_manager.sync_roles()
- num_pv_before = db.session().query(ab_perm_role).count()
+ num_pv_before =
db.session().scalars(select(func.count()).select_from(ab_perm_role)).one()
+
security_manager.sync_roles()
- num_pv_after = db.session().query(ab_perm_role).count()
+ num_pv_after =
db.session().scalars(select(func.count()).select_from(ab_perm_role)).one()
assert num_pv_before == num_pv_after
@@ -1119,8 +1120,8 @@ def test_permissions_work_for_dags_with_dot_in_dagname(
assert_user_has_dag_perms(perms=["GET", "PUT"], dag_id=dag_id,
user=user)
assert_user_does_not_have_dag_perms(perms=["GET", "PUT"],
dag_id=dag_id_2, user=user)
# Clean up DAG models and bundle
- session.query(DagModel).delete()
- session.query(DagBundleModel).filter(DagBundleModel.name ==
bundle_name).delete()
+ session.execute(delete(DagModel))
+ session.execute(delete(DagBundleModel).where(DagBundleModel.name
== bundle_name))
session.commit()
diff --git a/providers/fab/tests/unit/fab/utils.py
b/providers/fab/tests/unit/fab/utils.py
index 9d0afc60639..ad8f9791a6b 100644
--- a/providers/fab/tests/unit/fab/utils.py
+++ b/providers/fab/tests/unit/fab/utils.py
@@ -25,6 +25,7 @@ from unittest import mock
import flask
import pytest
+from sqlalchemy import delete, select
from airflow.models import Log
@@ -102,8 +103,8 @@ def _masked_value_check(data, sensitive_fields):
def _check_last_log(session, dag_id, event, logical_date, expected_extra=None,
check_masked=False):
- logs = (
- session.query(
+ logs = session.execute(
+ select(
Log.dag_id,
Log.task_id,
Log.event,
@@ -118,8 +119,7 @@ def _check_last_log(session, dag_id, event, logical_date,
expected_extra=None, c
)
.order_by(Log.dttm.desc())
.limit(5)
- .all()
- )
+ ).all()
assert len(logs) >= 1
assert logs[0].extra
if expected_extra:
@@ -128,28 +128,17 @@ def _check_last_log(session, dag_id, event, logical_date,
expected_extra=None, c
extra_json = json.loads(logs[0].extra)
_masked_value_check(extra_json, sensitive_fields)
- session.query(Log).delete()
+ session.execute(delete(Log))
def _check_last_log_masked_connection(session, dag_id, event, logical_date):
- logs = (
- session.query(
- Log.dag_id,
- Log.task_id,
- Log.event,
- Log.logical_date,
- Log.owner,
- Log.extra,
- )
- .filter(
- Log.dag_id == dag_id,
- Log.event == event,
- Log.logical_date == logical_date,
- )
+ logs = session.execute(
+ select(Log.dag_id, Log.task_id, Log.event, Log.logical_date,
Log.owner, Log.extra)
+ .where(Log.dag_id == dag_id, Log.event == event, Log.logical_date ==
logical_date)
.order_by(Log.dttm.desc())
.limit(5)
- .all()
- )
+ ).all()
+
assert len(logs) >= 1
extra = ast.literal_eval(logs[0].extra)
assert extra == {
@@ -170,8 +159,8 @@ def _check_last_log_masked_variable(
event,
logical_date,
):
- logs = (
- session.query(
+ logs = session.execute(
+ select(
Log.dag_id,
Log.task_id,
Log.event,
@@ -179,15 +168,14 @@ def _check_last_log_masked_variable(
Log.owner,
Log.extra,
)
- .filter(
+ .where(
Log.dag_id == dag_id,
Log.event == event,
Log.logical_date == logical_date,
)
.order_by(Log.dttm.desc())
.limit(5)
- .all()
- )
+ ).all()
assert len(logs) >= 1
extra_dict = ast.literal_eval(logs[0].extra)
assert extra_dict == {"key": "x_secret", "val": "***"}
diff --git
a/providers/fab/tests/unit/fab/www/views/test_views_custom_user_views.py
b/providers/fab/tests/unit/fab/www/views/test_views_custom_user_views.py
index cd018f2c109..e63ce584e3b 100644
--- a/providers/fab/tests/unit/fab/www/views/test_views_custom_user_views.py
+++ b/providers/fab/tests/unit/fab/www/views/test_views_custom_user_views.py
@@ -23,6 +23,7 @@ from unittest import mock
import pytest
from flask.sessions import SecureCookieSessionInterface
from flask_appbuilder import SQLA
+from sqlalchemy import delete, func, select
from airflow import settings
from airflow.providers.fab.www import app as application
@@ -214,7 +215,7 @@ class TestResetUserSessions:
self.model = self.interface.sql_session_model
self.serializer = self.interface.serializer
self.db = self.interface.db
- self.db.session.query(self.model).delete()
+ self.db.session.execute(delete(self.model))
self.db.session.commit()
self.db.session.flush()
self.user_1 = create_user(
@@ -256,7 +257,7 @@ class TestResetUserSessions:
self.create_user_db_session("session_id_2", time_delta, self.user_2.id)
self.db.session.commit()
self.db.session.flush()
- assert self.db.session.query(self.model).count() == 2
+ assert
self.db.session.scalar(select(func.count()).select_from(self.model)) == 2
assert self.get_session_by_id("session_id_1") is not None
assert self.get_session_by_id("session_id_2") is not None
@@ -264,14 +265,14 @@ class TestResetUserSessions:
self.db.session.commit()
self.db.session.flush()
if user_sessions_deleted:
- assert self.db.session.query(self.model).count() == 1
+ assert
self.db.session.scalar(select(func.count()).select_from(self.model)) == 1
assert self.get_session_by_id("session_id_1") is None
else:
- assert self.db.session.query(self.model).count() == 2
+ assert
self.db.session.scalar(select(func.count()).select_from(self.model)) == 2
assert self.get_session_by_id("session_id_1") is not None
def get_session_by_id(self, session_id: str):
- return self.db.session.query(self.model).filter(self.model.session_id
== session_id).scalar()
+ return
self.db.session.scalar(select(self.model).where(self.model.session_id ==
session_id))
@mock.patch("airflow.providers.fab.auth_manager.security_manager.override.flash")
@mock.patch(
@@ -285,7 +286,7 @@ class TestResetUserSessions:
self.create_user_db_session("session_id_2", timedelta(days=1),
self.user_2.id)
self.db.session.commit()
self.db.session.flush()
- assert self.db.session.query(self.model).count() == 2
+ assert
self.db.session.scalar(select(func.count()).select_from(self.model)) == 2
assert self.get_session_by_id("session_id_1") is not None
assert self.get_session_by_id("session_id_2") is not None
self.security_manager.reset_password(self.user_1.id, "new_password")
@@ -294,7 +295,7 @@ class TestResetUserSessions:
"The old sessions for user user_to_delete_1 have <b>NOT</b> been
deleted!"
in flash_mock.call_args[0][0]
)
- assert self.db.session.query(self.model).count() == 2
+ assert
self.db.session.scalar(select(func.count()).select_from(self.model)) == 2
assert self.get_session_by_id("session_id_1") is not None
assert self.get_session_by_id("session_id_2") is not None
@@ -320,7 +321,7 @@ class TestResetUserSessions:
self.create_user_db_session("session_id_2", timedelta(days=1),
self.user_2.id)
self.db.session.commit()
self.db.session.flush()
- assert self.db.session.query(self.model).count() == 2
+ assert
self.db.session.scalar(select(func.count()).select_from(self.model)) == 2
assert self.get_session_by_id("session_id_1") is not None
assert self.get_session_by_id("session_id_2") is not None
self.security_manager.reset_password(self.user_1.id, "new_password")
@@ -329,7 +330,7 @@ class TestResetUserSessions:
"The old sessions for user user_to_delete_1 have *NOT* been
deleted!\n"
in log_mock.warning.call_args[0][0]
)
- assert self.db.session.query(self.model).count() == 2
+ assert
self.db.session.scalar(select(func.count()).select_from(self.model)) == 2
assert self.get_session_by_id("session_id_1") is not None
assert self.get_session_by_id("session_id_2") is not None
