Migrate app code to SQLAlchemy 2.0 Query API

Replace all legacy session.query() and Model.query patterns with
session.execute(select()), session.scalars(), session.scalar(), and
db.paginate() across 10 files:

  flowapp/__init__.py, auth.py, messages.py
  flowapp/models/utils.py, rules/whitelist.py, community.py, user.py
  flowapp/services/base.py, rule_service.py, whitelist_service.py

Add test coverage for previously untested DB-touching code paths:
  tests/test_auth.py, test_model_utils.py, test_services_base.py
  (new files), plus additions to test_flowapp.py, test_messages.py,
  test_models.py, test_whitelist_service.py

Author: jirivrany

flowapp/__init__.py

@@ -2,6 +2,7 @@
 import os
 
 from flask import Flask, redirect, render_template, session, url_for, flash
+from sqlalchemy import select
 
 from flask_sso import SSO
 from flask_sqlalchemy import SQLAlchemy
@@ -128,7 +129,7 @@ def index():
     @auth_required
     def select_org(org_id=None):
         uuid = session.get("user_uuid")
-        user = db.session.query(models.User).filter_by(uuid=uuid).first()
+        user = db.session.execute(select(models.User).filter_by(uuid=uuid)).scalar_one_or_none()
 
         if user is None:
             return render_template("errors/404.html"), 404

flowapp/auth.py

@@ -1,6 +1,7 @@
 from functools import wraps
 from typing import List, Optional
 from flask import current_app, redirect, request, url_for, session, abort
+from sqlalchemy import select
 
 from flowapp import __version__, db, validators
 from flowapp.models import Flowspec4, Flowspec6, RTBH, Whitelist, get_user_nets
@@ -153,35 +154,35 @@ def get_user_allowed_rule_ids(rule_type: str, user_id: int, user_role_ids: List[
     # Admin users can modify any rules
     if 3 in user_role_ids:
         if rule_type == "ipv4":
-            return [r.id for r in db.session.query(Flowspec4.id).all()]
+            return list(db.session.scalars(select(Flowspec4.id)))
         elif rule_type == "ipv6":
-            return [r.id for r in db.session.query(Flowspec6.id).all()]
+            return list(db.session.scalars(select(Flowspec6.id)))
         elif rule_type == "rtbh":
-            return [r.id for r in db.session.query(RTBH.id).all()]
+            return list(db.session.scalars(select(RTBH.id)))
         elif rule_type == "whitelist":
-            return [r.id for r in db.session.query(Whitelist.id).all()]
+            return list(db.session.scalars(select(Whitelist.id)))
         return []
 
     # Regular users - filter by network ranges
     net_ranges = get_user_nets(user_id)
 
     if rule_type == "ipv4":
-        rules = db.session.query(Flowspec4).all()
+        rules = db.session.scalars(select(Flowspec4)).all()
         filtered_rules = validators.filter_rules_in_network(net_ranges, rules)
         return [r.id for r in filtered_rules]
 
     elif rule_type == "ipv6":
-        rules = db.session.query(Flowspec6).all()
+        rules = db.session.scalars(select(Flowspec6)).all()
         filtered_rules = validators.filter_rules_in_network(net_ranges, rules)
         return [r.id for r in filtered_rules]
 
     elif rule_type == "rtbh":
-        rules = db.session.query(RTBH).all()
+        rules = db.session.scalars(select(RTBH)).all()
         filtered_rules = validators.filter_rtbh_rules(net_ranges, rules)
         return [r.id for r in filtered_rules]
 
     elif rule_type == "whitelist":
-        rules = db.session.query(Whitelist).all()
+        rules = db.session.scalars(select(Whitelist)).all()
         filtered_rules = validators.filter_rules_in_network(net_ranges, rules)
         return [r.id for r in filtered_rules]
 

flowapp/messages.py

@@ -10,6 +10,7 @@
 )
 from flowapp.flowspec import translate_sequence as trps
 from flask import current_app
+from sqlalchemy import select
 from flowapp.models import ASPath
 from flowapp import db
 
@@ -138,7 +139,7 @@ def create_rtbh(rule, message_type=ANNOUNCE):
 
     as_path_string = ""
     if rule.community.as_path:
-        match = db.session.query(ASPath).filter(ASPath.prefix == source).first()
+        match = db.session.execute(select(ASPath).filter_by(prefix=source)).scalar_one_or_none()
         as_path_string = f"as-path [ {match.as_path} ]" if match else ""
 
     return "{neighbor}{action} route {source} next-hop {nexthop} {as_path} {community} {large_community} {extended_community}{rd_string}".format(

flowapp/models/community.py

@@ -1,4 +1,4 @@
-from sqlalchemy import event
+from sqlalchemy import event, select
 from .base import db
 
 
@@ -26,7 +26,7 @@ def __init__(self, name, comm, larcomm, extcomm, description, as_path, role_id):
 
     @classmethod
     def get_whitelistable_communities(cls, id_list):
-        return cls.query.filter(cls.id.in_(id_list)).all()
+        return db.session.scalars(select(cls).filter(cls.id.in_(id_list))).all()
 
     def __repr__(self):
         return f"<Community {self.name}>"

flowapp/models/rules/whitelist.py

@@ -1,6 +1,7 @@
 from flowapp import utils
 from ..base import db
 from datetime import datetime
+from sqlalchemy import func, select
 from flowapp.constants import RuleTypes, RuleOrigin
 
 
@@ -119,7 +120,7 @@ def get_by_whitelist_id(cls, whitelist_id: int):
         Returns:
             list: All RuleWhitelistCache objects with the specified whitelist_id
         """
-        return cls.query.filter_by(whitelist_id=whitelist_id).all()
+        return db.session.scalars(select(cls).filter_by(whitelist_id=whitelist_id)).all()
 
     @classmethod
     def clean_by_whitelist_id(cls, whitelist_id: int):
@@ -132,7 +133,9 @@ def clean_by_whitelist_id(cls, whitelist_id: int):
         Returns:
             int: Number of rows deleted
         """
-        deleted = cls.query.filter_by(whitelist_id=whitelist_id).delete()
+        deleted = db.session.execute(
+            db.delete(cls).filter_by(whitelist_id=whitelist_id)
+        ).rowcount
         db.session.commit()
         return deleted
 
@@ -147,7 +150,9 @@ def delete_by_rule_id(cls, rule_id: int):
         Returns:
             int: Number of rows deleted
         """
-        deleted = cls.query.filter_by(rid=rule_id).delete()
+        deleted = db.session.execute(
+            db.delete(cls).filter_by(rid=rule_id)
+        ).rowcount
         db.session.commit()
         return deleted
 
@@ -163,7 +168,9 @@ def count_by_rule(cls, rule_id: int, rule_type: RuleTypes):
         Returns:
             int: Number of cache entries
         """
-        return cls.query.filter_by(rid=rule_id, rtype=rule_type.value).count()
+        return db.session.scalar(
+            select(func.count()).select_from(cls).filter_by(rid=rule_id, rtype=rule_type.value)
+        )
 
     def __repr__(self):
         return f"<RuleWhitelistCache {self.rid} {self.rtype} {self.rorigin}>"

flowapp/models/user.py

@@ -1,4 +1,4 @@
-from sqlalchemy import event
+from sqlalchemy import event, select
 from .base import db, user_role, user_organization
 from .organization import Organization
 
@@ -46,12 +46,12 @@ def update(self, form):
             self.organization.remove(org)
 
         for role_id in form.role_ids.data:
-            my_role = db.session.query(Role).filter_by(id=role_id).first()
+            my_role = db.session.execute(select(Role).filter_by(id=role_id)).scalar_one()
             if my_role not in self.role:
                 self.role.append(my_role)
 
         for org_id in form.org_ids.data:
-            my_org = db.session.query(Organization).filter_by(id=org_id).first()
+            my_org = db.session.execute(select(Organization).filter_by(id=org_id)).scalar_one()
             if my_org not in self.organization:
                 self.organization.append(my_org)