aki-008
fix: enable query param auth for pdf viewer
1591012
from fastapi import Depends, HTTPException, status, Query
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select
from jose import JWTError, jwt
from app.database import async_session_maker
from app.models import User
from app.config import settings
from fastapi import Request
from chromadb import AsyncHttpClient
from chromadb.api.models.Collection import Collection
from typing import Optional
security = HTTPBearer(auto_error=False)
async def get_db():
async with async_session_maker() as session:
try:
yield session
await session.commit()
except Exception:
await session.rollback()
raise
finally:
await session.close()
async def get_current_user(
credentials: Optional[HTTPAuthorizationCredentials] = Depends(security),
token_query: Optional[str] = Query(None, alias="token"),
db: AsyncSession = Depends(get_db)
) -> User:
credentials_exception = HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Could not validate credentials",
headers={"WWW-Authenticate": "Bearer"},
)
token = None
if credentials:
token = credentials.credentials
elif token_query:
token = token_query
if not token:
raise credentials_exception
try:
payload = jwt.decode(token, settings.SECRET_KEY, algorithms=[settings.ALGORITHM])
username: str = payload.get("sub")
if username is None:
raise credentials_exception
except JWTError:
raise credentials_exception
result = await db.execute(select(User).filter(User.username == username))
user = result.scalar_one_or_none()
if user is None:
raise credentials_exception
return user
async def get_chroma_client(request: Request) -> AsyncHttpClient:
client = getattr(request.app.state, "chroma_client", None)
if client is None:
raise RuntimeError("ChromaDB client is not initialized in App State.")
return client
def get_chroma_collection(request: Request) -> Collection:
collection = getattr(request.app.state, "chroma_collection", None)
if collection is None:
raise RuntimeError("ChromaDB Collection not loaded during application startup.")
return collection