Auth middleware to avoid multiple queries

Author: kompotkot

brood/actions.py

@@ -5,7 +5,7 @@
 import logging
 from random import randint
 import re
-from typing import Any, cast, Callable, Dict, List, Optional, Set
+from typing import Any, cast, Callable, Dict, List, Optional, Set, Tuple
 import uuid
 
 from passlib.context import CryptContext
@@ -864,6 +864,63 @@ def get_token(session: Session, token: uuid.UUID) -> Token:
     return token_object
 
 
+def get_current_user_extended(
+    session: Session, token: uuid.UUID
+) -> Tuple[bool, data.UserExtendedResponse]:
+    """
+    Extend user data with groups for authentication middleware.
+    """
+    query = (
+        session.query(Token.active, User, Group, GroupUser)
+        .join(User, User.id == Token.user_id)
+        .join(GroupUser, GroupUser.user_id == User.id)
+        .join(Group, Group.id == GroupUser.group_id)
+        .filter(Token.id == token)
+        .filter(GroupUser.user_id == User.id)
+        .group_by(Token.active, User, Group, GroupUser)
+    )
+    objects = query.all()
+    if len(objects) == 0:
+        raise TokenNotFound(f"Token not found with ID: {token}")
+
+    active_token = objects[0][0]
+    user = objects[0][1]
+    groups = []
+    for object in objects:
+        if object[3].user_id != user.id:
+            logger.error(
+                f"Unexpected group id: {object[2].id} fetched for user with id: {user.id}"
+            )
+            raise Exception("Unexpected group in list")
+        groups.append(
+            data.GroupUserResponse(
+                group_id=object[2].id,
+                user_id=object[3].user_id,
+                user_type=object[3].user_type,
+                autogenerated=object[2].autogenerated,
+                group_name=object[2].name,
+                parent=object[2].parent,
+            )
+        )
+
+    user_extended = data.UserExtendedResponse(
+        user_id=user.id,
+        username=user.username,
+        first_name=user.first_name,
+        last_name=user.last_name,
+        email=user.email,
+        normalized_email=user.normalized_email,
+        verified=user.verified,
+        created_at=user.created_at,
+        updated_at=user.updated_at,
+        autogenerated=user.autogenerated,
+        application_id=user.application_id,
+        groups=groups,
+    )
+
+    return active_token, user_extended
+
+
 def update_token(
     session: Session,
     token: uuid.UUID,

brood/api.py

@@ -30,6 +30,7 @@
     oauth2_scheme,
     autogenerated_user_token_check,
     get_current_user,
+    current_user_extended,
     is_token_restricted,
     is_token_restricted_or_installation,
     get_current_user_or_installation,
@@ -461,6 +462,16 @@ async def get_user_by_id_handler(
     return user
 
 
+@app.get("/auth", tags=["users"], response_model=data.UserExtendedResponse)
+async def get_auth_handler(
+    current_user: data.UserExtendedResponse = Depends(current_user_extended),
+) -> data.UserExtendedResponse:
+    """
+    Authorization middleware.
+    """
+    return current_user
+
+
 @app.post("/confirm", tags=["users"], response_model=data.UserResponse)
 async def verification_handler(
     token_restricted: bool = Depends(is_token_restricted),

brood/data.py

@@ -273,3 +273,7 @@ class ApplicationResponse(BaseModel):
 
 class ApplicationsListResponse(BaseModel):
     applications: List[ApplicationResponse] = Field(default_factory=list)
+
+
+class UserExtendedResponse(UserResponse):
+    groups: List[GroupUserResponse] = Field(default_factory=list)

brood/middleware.py

@@ -9,6 +9,7 @@
 from fastapi.security import OAuth2PasswordBearer
 
 from . import actions
+from . import data
 from . import models
 from .external import yield_db_session_from_env
 from .settings import BOT_INSTALLATION_TOKEN, BOT_INSTALLATION_TOKEN_HEADER
@@ -32,6 +33,24 @@ async def get_current_user(
     return token_object.user
 
 
+async def current_user_extended(
+    token: UUID = Depends(oauth2_scheme),
+    db_session=Depends(yield_db_session_from_env),
+) -> data.UserExtendedResponse:
+    try:
+        token_active, user_extended = actions.get_current_user_extended(
+            session=db_session, token=token
+        )
+    except actions.TokenNotFound:
+        raise HTTPException(status_code=404, detail="Access token not found")
+    except Exception:
+        raise HTTPException(status_code=500)
+    if not token_active:
+        raise HTTPException(status_code=403, detail="Token has expired")
+
+    return user_extended
+
+
 def autogenerated_user_token_check(request: Request) -> bool:
     if BOT_INSTALLATION_TOKEN is None:
         raise ValueError("BOT_INSTALLATION_TOKEN environment variable must be set")