diff --git a/src/govoplan_access/backend/administration.py b/src/govoplan_access/backend/administration.py index 1683ecc..60aeea0 100644 --- a/src/govoplan_access/backend/administration.py +++ b/src/govoplan_access/backend/administration.py @@ -42,6 +42,69 @@ class SqlAccessAdministration(AccessAdministration): "active_api_keys": int(active_api_keys), } + def tenant_counts_many( + self, + session: object, + tenant_ids: Sequence[str], + ) -> Mapping[str, Mapping[str, int]]: + ids = tuple(dict.fromkeys(str(tenant_id) for tenant_id in tenant_ids if tenant_id)) + if not ids: + return {} + db = _session(session) + counts: dict[str, dict[str, int]] = { + tenant_id: { + "users": 0, + "active_users": 0, + "groups": 0, + "api_keys": 0, + "active_api_keys": 0, + } + for tenant_id in ids + } + user_rows = ( + db.query( + User.tenant_id, + func.count(User.id), + func.coalesce( + func.sum(case((User.is_active.is_(True), 1), else_=0)), + 0, + ), + ) + .filter(User.tenant_id.in_(ids)) + .group_by(User.tenant_id) + .all() + ) + for tenant_id, users, active_users in user_rows: + counts[tenant_id]["users"] = int(users) + counts[tenant_id]["active_users"] = int(active_users) + + group_rows = ( + db.query(Group.tenant_id, func.count(Group.id)) + .filter(Group.tenant_id.in_(ids)) + .group_by(Group.tenant_id) + .all() + ) + for tenant_id, groups in group_rows: + counts[tenant_id]["groups"] = int(groups) + + api_key_rows = ( + db.query( + ApiKey.tenant_id, + func.count(ApiKey.id), + func.coalesce( + func.sum(case((ApiKey.revoked_at.is_(None), 1), else_=0)), + 0, + ), + ) + .filter(ApiKey.tenant_id.in_(ids)) + .group_by(ApiKey.tenant_id) + .all() + ) + for tenant_id, api_keys, active_api_keys in api_key_rows: + counts[tenant_id]["api_keys"] = int(api_keys) + counts[tenant_id]["active_api_keys"] = int(active_api_keys) + return counts + def system_account_count(self, session: object) -> int: db = _session(session) return db.query(Account).count() diff --git a/tests/test_admin_batch_helpers.py b/tests/test_admin_batch_helpers.py index c14a627..b1d83e7 100644 --- a/tests/test_admin_batch_helpers.py +++ b/tests/test_admin_batch_helpers.py @@ -1,10 +1,12 @@ from __future__ import annotations import unittest +from datetime import UTC, datetime -from sqlalchemy import create_engine +from sqlalchemy import create_engine, event from sqlalchemy.orm import sessionmaker +from govoplan_access.backend.administration import SqlAccessAdministration from govoplan_access.backend.api.v1.admin_common import ( _accounts_by_user_id, _group_member_ids_by_group_id, @@ -13,7 +15,7 @@ from govoplan_access.backend.api.v1.admin_common import ( _roles_by_user_id, _tenant_role_assignment_counts, ) -from govoplan_access.backend.db.models import Account, Group, GroupRoleAssignment, Role, User, UserGroupMembership, UserRoleAssignment +from govoplan_access.backend.db.models import Account, ApiKey, Group, GroupRoleAssignment, Role, User, UserGroupMembership, UserRoleAssignment from govoplan_core.db.base import Base @@ -56,6 +58,92 @@ class AdminBatchHelperTests(unittest.TestCase): self.assertEqual([item.id for item in roles_by_user[user.id]], [role.id]) self.assertEqual(role_counts, {role.id: (1, 1)}) + def test_tenant_counts_many_uses_three_grouped_queries(self) -> None: + accounts = [ + Account( + id=f"account-{index}", + email=f"user-{index}@example.test", + normalized_email=f"user-{index}@example.test", + ) + for index in range(3) + ] + users = [ + User( + id="user-1", + tenant_id="tenant-1", + account_id=accounts[0].id, + email=accounts[0].email, + ), + User( + id="user-2", + tenant_id="tenant-1", + account_id=accounts[1].id, + email=accounts[1].email, + is_active=False, + ), + User( + id="user-3", + tenant_id="tenant-2", + account_id=accounts[2].id, + email=accounts[2].email, + ), + ] + self.session.add_all( + [ + *accounts, + *users, + Group(id="group-1", tenant_id="tenant-1", slug="one", name="One"), + Group(id="group-2", tenant_id="tenant-2", slug="two", name="Two"), + ApiKey( + id="key-1", + tenant_id="tenant-1", + user_id="user-1", + name="Active", + prefix="active", + key_hash="hash-1", + ), + ApiKey( + id="key-2", + tenant_id="tenant-1", + user_id="user-2", + name="Revoked", + prefix="revoked", + key_hash="hash-2", + revoked_at=datetime.now(UTC), + ), + ] + ) + self.session.commit() + + query_count = 0 + + def count_query(*_args: object) -> None: + nonlocal query_count + query_count += 1 + + event.listen(self.engine, "before_cursor_execute", count_query) + try: + counts = SqlAccessAdministration().tenant_counts_many( + self.session, + ["tenant-1", "tenant-2", "tenant-empty"], + ) + finally: + event.remove(self.engine, "before_cursor_execute", count_query) + + self.assertEqual(3, query_count) + self.assertEqual( + { + "users": 2, + "active_users": 1, + "groups": 1, + "api_keys": 2, + "active_api_keys": 1, + }, + counts["tenant-1"], + ) + self.assertEqual(1, counts["tenant-2"]["users"]) + self.assertEqual(0, counts["tenant-empty"]["users"]) + if __name__ == "__main__": unittest.main()