fix(authentication): enhance admin retrieval with metrics and refactor related functions

Author: x0sina

app/routers/admin.py

@@ -23,7 +23,13 @@
 from app.utils import responses
 from app.utils.jwt import create_admin_token
 
-from .authentication import check_sudo_admin, get_current, validate_admin, validate_mini_app_admin
+from .authentication import (
+    check_sudo_admin,
+    get_current,
+    get_current_with_metrics,
+    validate_admin,
+    validate_mini_app_admin,
+)
 
 router = APIRouter(tags=["Admin"], prefix="/api/admin", responses={401: responses._401, 403: responses._403})
 admin_operator = AdminOperation(operator_type=OperatorType.API)
@@ -130,7 +136,7 @@ async def remove_admin(
 
 
 @router.get("", response_model=AdminDetails)
-def get_current_admin(admin: AdminDetails = Depends(get_current)):
+def get_current_admin(admin: AdminDetails = Depends(get_current_with_metrics)):
     """Retrieve the current authenticated admin."""
     return admin
 

app/routers/authentication.py

@@ -4,9 +4,11 @@
 from fastapi import Depends, HTTPException, status
 from fastapi.security import OAuth2PasswordBearer
 
+from sqlalchemy import func, select
 
 from app.db import AsyncSession, get_db
 from app.db.crud.admin import find_admins_by_telegram_id, get_admin as get_admin_by_username, get_admin_by_telegram_id
+from app.db.models import Admin, AdminUsageLogs, User
 from app.models.admin import AdminDetails, AdminValidationResult, verify_password
 from app.models.settings import Telegram
 from app.settings import telegram_settings
@@ -16,20 +18,83 @@
 oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/api/admin/token")
 
 
+def _build_admin_details(
+    db_admin: Admin,
+    *,
+    total_users: int = 0,
+    reseted_usage: int | None = None,
+) -> AdminDetails:
+    used_traffic = int(db_admin.used_traffic or 0)
+    return AdminDetails(
+        id=db_admin.id,
+        username=db_admin.username,
+        is_sudo=db_admin.is_sudo,
+        total_users=int(total_users or 0),
+        used_traffic=used_traffic,
+        is_disabled=db_admin.is_disabled,
+        telegram_id=db_admin.telegram_id,
+        discord_webhook=db_admin.discord_webhook,
+        sub_domain=db_admin.sub_domain,
+        profile_title=db_admin.profile_title,
+        support_url=db_admin.support_url,
+        note=db_admin.note,
+        notification_enable=db_admin.notification_enable,
+        discord_id=db_admin.discord_id,
+        sub_template=db_admin.sub_template,
+        lifetime_used_traffic=None if reseted_usage is None else int(reseted_usage or 0) + used_traffic,
+    )
+
+
+def _is_token_valid_for_admin(db_admin: Admin, payload: dict) -> bool:
+    if not db_admin.password_reset_at:
+        return True
+    if not payload.get("created_at"):
+        return False
+    return db_admin.password_reset_at.astimezone(tz.utc) <= payload.get("created_at")
+
+
 async def get_admin(db: AsyncSession, token: str) -> AdminDetails | None:
     payload = await get_admin_payload(token)
     if not payload:
         return
 
-    db_admin = await get_admin_by_username(db, payload["username"], load_users=True, load_usage_logs=True)
+    db_admin = await get_admin_by_username(db, payload["username"], load_users=False, load_usage_logs=False)
     if db_admin:
-        if db_admin.password_reset_at:
-            if not payload.get("created_at"):
-                return
-            if db_admin.password_reset_at.astimezone(tz.utc) > payload.get("created_at"):
-                return
+        if not _is_token_valid_for_admin(db_admin, payload):
+            return
 
-        return AdminDetails.model_validate(db_admin)
+        return _build_admin_details(db_admin)
+
+    elif payload["username"] in SUDOERS and payload["is_sudo"] is True:
+        return AdminDetails(username=payload["username"], is_sudo=True)
+
+
+async def get_admin_with_metrics(db: AsyncSession, token: str) -> AdminDetails | None:
+    payload = await get_admin_payload(token)
+    if not payload:
+        return
+
+    total_users_subquery = (
+        select(func.count(User.id)).where(User.admin_id == Admin.id).correlate(Admin).scalar_subquery()
+    )
+    reseted_usage_subquery = (
+        select(func.coalesce(func.sum(AdminUsageLogs.used_traffic_at_reset), 0))
+        .where(AdminUsageLogs.admin_id == Admin.id)
+        .correlate(Admin)
+        .scalar_subquery()
+    )
+    admin_row = (
+        await db.execute(
+            select(Admin, total_users_subquery, reseted_usage_subquery).where(Admin.username == payload["username"])
+        )
+    ).one_or_none()
+
+    if admin_row:
+        db_admin, total_users, reseted_usage = admin_row
+        if not _is_token_valid_for_admin(db_admin, payload):
+            return
+
+        return _build_admin_details(db_admin, total_users=total_users, reseted_usage=reseted_usage)
 
     elif payload["username"] in SUDOERS and payload["is_sudo"] is True:
         return AdminDetails(username=payload["username"], is_sudo=True)
@@ -53,6 +118,24 @@ async def get_current(db: AsyncSession = Depends(get_db), token: str = Depends(o
     return admin
 
 
+async def get_current_with_metrics(db: AsyncSession = Depends(get_db), token: str = Depends(oauth2_scheme)):
+    admin: AdminDetails | None = await get_admin_with_metrics(db, token)
+    if not admin:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Could not validate credentials",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+    if admin.is_disabled:
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail="your account has been disabled",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+
+    return admin
+
+
 async def check_sudo_admin(admin: AdminDetails = Depends(get_current)):
     if not admin.is_sudo:
         raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="You're not allowed")

tests/api/test_admin.py

@@ -9,7 +9,7 @@
 from sqlalchemy import select
 
 from app.db.crud.admin import get_admin_by_telegram_id
-from app.db.models import Admin, NodeUserUsage
+from app.db.models import Admin, AdminUsageLogs, NodeUserUsage
 from app.models.settings import RunMethod, Telegram
 from app.routers.authentication import validate_mini_app_admin
 from tests.api import TestSession, client
@@ -93,6 +93,69 @@ def test_get_admin(access_token):
     assert response.json()["username"] == username
 
 
+def test_get_admin_uses_aggregate_metrics_without_loading_relationships(access_token, monkeypatch: pytest.MonkeyPatch):
+    admin = create_admin(access_token)
+    admin_token_response = client.post(
+        url="/api/admin/token",
+        data={"username": admin["username"], "password": admin["password"], "grant_type": "password"},
+    )
+    assert admin_token_response.status_code == status.HTTP_200_OK
+    admin_token = admin_token_response.json()["access_token"]
+
+    user = create_user(admin_token, payload={"username": unique_name("admin_metric_user")})
+
+    async def _seed_admin_usage():
+        async with TestSession() as session:
+            result = await session.execute(select(Admin).where(Admin.username == admin["username"]))
+            db_admin = result.scalar_one()
+            db_admin.used_traffic = 12345
+            session.add(AdminUsageLogs(admin_id=db_admin.id, used_traffic_at_reset=6789))
+            await session.commit()
+
+    async def _assert_lightweight_admin_load(_, load_users: bool = True, load_usage_logs: bool = True):
+        assert load_users is False
+        assert load_usage_logs is False
+
+    try:
+        asyncio.run(_seed_admin_usage())
+        with monkeypatch.context() as patch_context:
+            patch_context.setattr("app.db.crud.admin.load_admin_attrs", _assert_lightweight_admin_load)
+            response = client.get(url="/api/admin", headers=auth_headers(admin_token))
+
+        assert response.status_code == status.HTTP_200_OK
+        data = response.json()
+        assert data["username"] == admin["username"]
+        assert data["total_users"] == 1
+        assert data["used_traffic"] == 12345
+        assert data["lifetime_used_traffic"] == 19134
+    finally:
+        delete_user(admin_token, user["username"])
+        delete_admin(access_token, admin["username"])
+
+
+def test_protected_routes_use_lightweight_current_admin(access_token, monkeypatch: pytest.MonkeyPatch):
+    admin = create_admin(access_token)
+    admin_token_response = client.post(
+        url="/api/admin/token",
+        data={"username": admin["username"], "password": admin["password"], "grant_type": "password"},
+    )
+    assert admin_token_response.status_code == status.HTTP_200_OK
+    admin_token = admin_token_response.json()["access_token"]
+
+    async def _assert_lightweight_admin_load(_, load_users: bool = True, load_usage_logs: bool = True):
+        assert load_users is False
+        assert load_usage_logs is False
+
+    try:
+        with monkeypatch.context() as patch_context:
+            patch_context.setattr("app.db.crud.admin.load_admin_attrs", _assert_lightweight_admin_load)
+            response = client.get(url="/api/users", headers=auth_headers(admin_token))
+
+        assert response.status_code == status.HTTP_200_OK
+    finally:
+        delete_admin(access_token, admin["username"])
+
+
 def test_admin_create(access_token):
     """Test that the admin create route is accessible."""