Migrate views to SQLAlchemy 2.0 Query API with SoC refactoring

Replace all legacy `session.query()` / `Model.query` calls in views
(api_common, api_keys, dashboard, rules, admin) with SQLAlchemy 2.0
`select()` API. Move DB queries out of views into model classmethods
and a new `get_org_rule_stats()` utility, following separation of
concerns. Add tests for all new classmethods before implementation.

New classmethods: ApiKey/MachineApiKey.get_by_key/user_id/all,
Flowspec4/6/RTBH.get_all_ordered/count_active, Community.get_all,
Action.get_all/get_all_ordered, RuleWhitelistCache.get_by_rule_ids,
User.get_all/get_all_ordered/get_by_uuid, Role/Organization classmethods,
ASPath.get_all/get_by_prefix, Log.get_recent_paginated.

Author: jirivrany

flowapp/models/api.py

@@ -1,4 +1,6 @@
 from datetime import datetime
+from typing import Optional
+from sqlalchemy import select
 from .base import db
 
 
@@ -20,6 +22,14 @@ def is_expired(self):
         else:
             return self.expires < datetime.now()
 
+    @classmethod
+    def get_by_key(cls, key: str) -> Optional["ApiKey"]:
+        return db.session.scalars(select(cls).filter_by(key=key)).first()
+
+    @classmethod
+    def get_by_user_id(cls, user_id: int) -> list:
+        return db.session.scalars(select(cls).filter_by(user_id=user_id)).all()
+
 
 class MachineApiKey(db.Model):
     id = db.Column(db.Integer, primary_key=True)
@@ -38,3 +48,11 @@ def is_expired(self):
             return False  # Non-expiring key
         else:
             return self.expires < datetime.now()
+
+    @classmethod
+    def get_by_key(cls, key: str) -> Optional["MachineApiKey"]:
+        return db.session.scalars(select(cls).filter_by(key=key)).first()
+
+    @classmethod
+    def get_all(cls) -> list:
+        return db.session.scalars(select(cls)).all()

flowapp/models/community.py

@@ -1,3 +1,4 @@
+from typing import List, Optional
 from sqlalchemy import event, select
 from .base import db
 
@@ -24,6 +25,10 @@ def __init__(self, name, comm, larcomm, extcomm, description, as_path, role_id):
         self.as_path = as_path
         self.role_id = role_id
 
+    @classmethod
+    def get_all(cls) -> List["Community"]:
+        return db.session.scalars(select(cls)).all()
+
     @classmethod
     def get_whitelistable_communities(cls, id_list):
         return db.session.scalars(select(cls).filter(cls.id.in_(id_list))).all()
@@ -42,7 +47,13 @@ class ASPath(db.Model):
     prefix = db.Column(db.String(120), unique=True)
     as_path = db.Column(db.String(250))
 
-    # Methods and initializer
+    @classmethod
+    def get_all(cls) -> List["ASPath"]:
+        return db.session.scalars(select(cls)).all()
+
+    @classmethod
+    def get_by_prefix(cls, prefix: str) -> Optional["ASPath"]:
+        return db.session.scalars(select(cls).filter_by(prefix=prefix)).first()
 
 
 # Note: seed data is also defined in migrations/versions/001_baseline.py - keep in sync

flowapp/models/log.py

@@ -1,5 +1,6 @@
 from datetime import datetime, timedelta
 
+from sqlalchemy import select
 from flowapp.constants import RuleTypes
 from .base import db
 
@@ -23,6 +24,17 @@ def __init__(self, time, task, user_id, rule_type, rule_id, author):
         self.user_id = user_id
         self.author = author
 
+    @classmethod
+    def get_recent_paginated(cls, page: int, per_page: int = 20, weeks: int = 1):
+        since = datetime.now() - timedelta(weeks=weeks)
+        return db.paginate(
+            select(cls).filter(cls.time > since).order_by(cls.time.desc()),
+            page=page,
+            per_page=per_page,
+            max_per_page=None,
+            error_out=False,
+        )
+
     @classmethod
     def delete_old(cls, days: int = 30):
         """Delete logs older than :param days from the database"""

flowapp/models/organization.py

@@ -1,4 +1,5 @@
-from sqlalchemy import event
+from typing import List, Optional
+from sqlalchemy import event, select
 from .base import db
 
 
@@ -27,6 +28,14 @@ def get_users(self):
         # self.user is the backref from the user_organization relationship
         return self.user
 
+    @classmethod
+    def get_all_ordered(cls) -> List["Organization"]:
+        return db.session.scalars(select(cls).order_by(cls.name)).all()
+
+    @classmethod
+    def get_by_name(cls, name: str) -> Optional["Organization"]:
+        return db.session.scalars(select(cls).filter_by(name=name)).first()
+
 
 # Event listeners for Organization
 # Note: seed data is also defined in migrations/versions/001_baseline.py - keep in sync

flowapp/models/rules/base.py

@@ -1,4 +1,5 @@
-from sqlalchemy import event
+from typing import List
+from sqlalchemy import event, select
 from ..base import db
 
 
@@ -30,6 +31,14 @@ def __init__(self, name, command, description, role_id=2):
         self.description = description
         self.role_id = role_id
 
+    @classmethod
+    def get_all_ordered(cls) -> List["Action"]:
+        return db.session.scalars(select(cls).order_by(cls.name)).all()
+
+    @classmethod
+    def get_all(cls) -> List["Action"]:
+        return db.session.scalars(select(cls)).all()
+
 
 # Event listeners for Rstate
 # Note: seed data is also defined in migrations/versions/001_baseline.py - keep in sync

flowapp/models/rules/flowspec.py

@@ -1,6 +1,8 @@
 import json
 from datetime import datetime
+from typing import List, Optional
 from flowapp import utils
+from sqlalchemy import func, select
 from ..base import db
 
 
@@ -163,6 +165,17 @@ def json(self, prefered_format="yearfirst"):
         """
         return json.dumps(self.to_dict())
 
+    @classmethod
+    def get_all_ordered(cls) -> List["Flowspec4"]:
+        return db.session.scalars(select(cls).order_by(cls.expires.desc())).all()
+
+    @classmethod
+    def count_active(cls, org_id: Optional[int] = None) -> int:
+        q = select(func.count()).select_from(cls).filter_by(rstate_id=1)
+        if org_id is not None:
+            q = select(func.count()).select_from(cls).filter_by(rstate_id=1, org_id=org_id)
+        return db.session.scalar(q)
+
 
 class Flowspec6(db.Model):
     id = db.Column(db.Integer, primary_key=True)
@@ -292,3 +305,14 @@ def json(self, prefered_format="yearfirst"):
         :returns: json
         """
         return json.dumps(self.to_dict())
+
+    @classmethod
+    def get_all_ordered(cls) -> List["Flowspec6"]:
+        return db.session.scalars(select(cls).order_by(cls.expires.desc())).all()
+
+    @classmethod
+    def count_active(cls, org_id: Optional[int] = None) -> int:
+        q = select(func.count()).select_from(cls).filter_by(rstate_id=1)
+        if org_id is not None:
+            q = select(func.count()).select_from(cls).filter_by(rstate_id=1, org_id=org_id)
+        return db.session.scalar(q)

flowapp/models/rules/rtbh.py

@@ -1,6 +1,8 @@
 import json
 from datetime import datetime
+from typing import List, Optional
 from flowapp import utils
+from sqlalchemy import func, select
 from ..base import db
 
 
@@ -151,3 +153,14 @@ def __str__(self):
 
     def get_author(self):
         return f"{self.user.email} / {self.org}"
+
+    @classmethod
+    def get_all_ordered(cls) -> List["RTBH"]:
+        return db.session.scalars(select(cls).order_by(cls.expires.desc())).all()
+
+    @classmethod
+    def count_active(cls, org_id: Optional[int] = None) -> int:
+        q = select(func.count()).select_from(cls).filter_by(rstate_id=1)
+        if org_id is not None:
+            q = select(func.count()).select_from(cls).filter_by(rstate_id=1, org_id=org_id)
+        return db.session.scalar(q)

flowapp/models/rules/whitelist.py

@@ -156,6 +156,14 @@ def delete_by_rule_id(cls, rule_id: int):
         db.session.commit()
         return deleted
 
+    @classmethod
+    def get_by_rule_ids(cls, rule_ids: list, rule_type: "RuleTypes") -> list:
+        if not rule_ids:
+            return []
+        return db.session.scalars(
+            select(cls).filter(cls.rid.in_(rule_ids), cls.rtype == rule_type.value)
+        ).all()
+
     @classmethod
     def count_by_rule(cls, rule_id: int, rule_type: RuleTypes):
         """

flowapp/models/user.py

@@ -1,3 +1,4 @@
+from typing import List, Optional
 from sqlalchemy import event, select
 from .base import db, user_role, user_organization
 from .organization import Organization
@@ -57,6 +58,18 @@ def update(self, form):
 
         db.session.commit()
 
+    @classmethod
+    def get_all(cls) -> List["User"]:
+        return db.session.scalars(select(cls)).all()
+
+    @classmethod
+    def get_all_ordered(cls) -> List["User"]:
+        return db.session.scalars(select(cls).order_by(cls.name)).all()
+
+    @classmethod
+    def get_by_uuid(cls, uuid: str) -> Optional["User"]:
+        return db.session.scalars(select(cls).filter_by(uuid=uuid)).first()
+
 
 class Role(db.Model):
     id = db.Column(db.Integer, primary_key=True)
@@ -70,6 +83,10 @@ def __init__(self, name, description):
     def __repr__(self):
         return self.name
 
+    @classmethod
+    def get_all_ordered(cls) -> List["Role"]:
+        return db.session.scalars(select(cls).order_by(cls.name)).all()
+
 
 # Event listeners for Role
 # Note: seed data is also defined in migrations/versions/001_baseline.py - keep in sync

flowapp/models/utils.py

@@ -65,6 +65,47 @@ def check_global_rule_limit(rule_type: RuleTypes) -> bool:
     return False
 
 
+def get_org_rule_stats() -> dict:
+    """
+    Return per-org and global active rule counts for the admin organizations view.
+    """
+    orgs = db.session.scalars(
+        select(Organization).options(db.joinedload(Organization.rtbh))
+    ).unique().all()
+
+    rtbh_counts = dict(
+        db.session.execute(
+            select(RTBH.org_id, func.count(RTBH.id))
+            .filter(RTBH.rstate_id == 1)
+            .group_by(RTBH.org_id)
+        ).all()
+    )
+    flowspec4_counts = dict(
+        db.session.execute(
+            select(Flowspec4.org_id, func.count(Flowspec4.id))
+            .filter(Flowspec4.rstate_id == 1)
+            .group_by(Flowspec4.org_id)
+        ).all()
+    )
+    flowspec6_counts = dict(
+        db.session.execute(
+            select(Flowspec6.org_id, func.count(Flowspec6.id))
+            .filter(Flowspec6.rstate_id == 1)
+            .group_by(Flowspec6.org_id)
+        ).all()
+    )
+
+    return {
+        "orgs": orgs,
+        "rtbh_counts": rtbh_counts,
+        "flowspec4_counts": flowspec4_counts,
+        "flowspec6_counts": flowspec6_counts,
+        "rtbh_all_count": RTBH.count_active(),
+        "flowspec4_all_count": Flowspec4.count_active(),
+        "flowspec6_all_count": Flowspec6.count_active(),
+    }
+
+
 def get_whitelist_model_if_exists(form_data):
     """
     Check if the record in database exist