remove session.query from fab/src

Author: Prab-27

providers/fab/src/airflow/providers/fab/auth_manager/security_manager/override.py

@@ -545,7 +545,8 @@ def reset_user_sessions(self, user: User) -> None:
             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.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(
@@ -560,7 +561,8 @@ def reset_user_sessions(self, user: User) -> None:
                     "warning",
                 )
             else:
-                for s in session.query(user_session_model):
+                # for s in session.query(user_session_model):
+                for s in session.scalars(user_session_model).all():
                     session_details = interface.serializer.loads(want_bytes(s.data))
                     if session_details.get("_user_id") == user.id:
                         session.delete(s)
@@ -1274,10 +1276,11 @@ def find_role(self, name):
 
         :param name: the role name
         """
-        return self.get_session.query(self.role_model).filter_by(name=name).one_or_none()
+        return self.get_session.execute(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.query(self.role_model).all()
+        return self.get_session.scalars(select(self.role_model)).all()
 
     def delete_role(self, role_name: str) -> None:
         """
@@ -1286,7 +1289,8 @@ def delete_role(self, role_name: str) -> None:
         :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.query(Role).filter(Role.name == role_name).first()
+        role = session.execute(select(Role).where(Role.name == role_name)).first()
         if role:
             log.info("Deleting role '%s'", role_name)
             session.delete(role)
@@ -1320,7 +1324,10 @@ def get_roles_from_keys(self, role_keys: list[str]) -> set[Role]:
         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.query(self.role_model).filter_by(name=self.auth_role_public).one_or_none()
+        return self.get_session.execute(
+            select(self.role_model).filter_by(name=self.auth_role_public)
+        ).one_or_none()
 
     """
     -----------
@@ -1377,7 +1384,8 @@ def get_user_by_id(self, pk):
 
     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.query(func.count(self.user_model.id)).scalar()
+        return self.get_session.execute(select(func.count(self.user_model.id))).scalar()
 
     def add_register_user(self, username, first_name, last_name, email, password="", hashed_password=""):
         """
@@ -1409,22 +1417,32 @@ def find_user(self, username=None, email=None):
         if username:
             try:
                 if self.auth_username_ci:
-                    return (
+                    """ return (
                         self.get_session.query(self.user_model)
                         .filter(func.lower(self.user_model.username) == func.lower(username))
                         .one_or_none()
-                    )
-                return (
+                    ) """
+                    return self.get_session.execute(
+                        select(self.user_model).where(
+                            func.lower(self.user_model.username) == func.lower(username)
+                        )
+                    ).one_or_none()
+                """ 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.execute(
+                    select(self.user_model).where(
+                        func.lower(self.user_model.username) == func.lower(username)
+                    )
+                ).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.execute(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
@@ -1456,7 +1474,7 @@ def del_register_user(self, register_user):
             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):
         """
@@ -1496,7 +1514,7 @@ def get_action(self, name: str) -> Action:
 
         :param name: name
         """
-        return self.get_session.query(self.action_model).filter_by(name=name).one_or_none()
+        return self.get_session.execute(select(self.action_model).filter_by(name=name)).one_or_none()
 
     def create_action(self, name):
         """
@@ -1529,11 +1547,9 @@ def delete_action(self, name: str) -> bool:
             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
@@ -1557,7 +1573,7 @@ def get_resource(self, name: str) -> Resource | None:
 
         :param name: Name of resource
         """
-        return self.get_session.query(self.resource_model).filter_by(name=name).one_or_none()
+        return self.get_session.execute(select(self.resource_model).filter_by(name=name)).one_or_none()
 
     def create_resource(self, name) -> Resource | None:
         """
@@ -1598,11 +1614,9 @@ def get_permission(
         action = self.get_action(action_name)
         resource = self.get_resource(resource_name)
         if action and resource:
-            return (
-                self.get_session.query(self.permission_model)
-                .filter_by(action=action, resource=resource)
-                .one_or_none()
-            )
+            return self.get_session.execute(
+                select(self.permission_model).filter_by(action=action, resource=resource)
+            ).one_or_none()
         return None
 
     def get_resource_permissions(self, resource: Resource) -> Permission:
@@ -1611,7 +1625,9 @@ def get_resource_permissions(self, resource: Resource) -> Permission:
 
         :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:
         """
@@ -1658,9 +1674,9 @@ def delete_permission(self, action_name: str, resource_name: str) -> None:
         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.execute(
+            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
@@ -1669,7 +1685,9 @@ def delete_permission(self, action_name: str, resource_name: str) -> None:
             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: