[python fastapi] Depends에 동적으로 parameter 사용하기

Why?

코드의 깔끔함을 위해서, 조금 더 재사용성이 좋고, API에 명시적으로 권한을 보고 싶어서 진행.

 

본문

fastapi에 Depends()를 이용해서, user를 판별하고 있었고, 특정 API 마다 접근 할 수 있는 role을 부여하고 싶었다.

예를 들면, admin만 접근 할 수 있는 API라던지, 로그인 한 유저라면 얼마든지 사용할 수 있는 API 라던지.

 

다만, 고려했던 사항은, router 내부에 코드를 추가하고 싶지 않았다. router의 Depends() 레벨에서 해결하고 싶었다.

 

아래는 간단하게 작성한, prototype code 이다.

from fastapi import HTTPException, Depends, FastAPI
from fastapi.security import OAuth2PasswordBearer
from jose import JWTError, jwt
from pydantic import BaseModel
from sqlalchemy import create_engine, Column, String, Integer
from sqlalchemy.orm import sessionmaker, Session, declarative_base
from enum import Enum
import uvicorn

# SQLite DB 설정
DATABASE_URL = "sqlite:///./test.db"

engine = create_engine(DATABASE_URL)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()

# OAuth2 토큰 스킴 정의
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")


# 유저 모델 정의
class User(BaseModel):
    username: str
    role: str


# SQLAlchemy UserModel 정의
class UserModel(Base):
    __tablename__ = "users"

    id = Column(Integer, primary_key=True, index=True)
    username = Column(String, unique=True, index=True)
    role = Column(String)


# 데이터베이스 초기화 (테이블 생성)
Base.metadata.create_all(bind=engine)


def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


def verify_jwt_token(token: str = Depends(oauth2_scheme)) -> User:
    try:
        payload = jwt.decode(token, "SECRET_KEY", algorithms=["HS256"])
        username: str = payload.get("sub")
        role: str = payload.get("role")

        return User(username=username, role=role)

    except JWTError:
        raise HTTPException(status_code=401, detail="Invalid token")


class RoleChecker:
    def __init__(self, allowed_roles):
        self.allowed_roles = allowed_roles

    def __call__(
        self, db: Session = Depends(get_db), user: User = Depends(verify_jwt_token)
    ):
        db_user = (
            db.query(UserModel).filter(UserModel.username == user.username).first()
        )
        if not db_user:
            raise HTTPException(status_code=404, detail="User not found")

        # 유저의 역할이 요구되는 역할을 포함하고 있는지 확인
        if db_user.role not in self.allowed_roles:
            raise HTTPException(status_code=403, detail="Permission denied")

        return db_user


app = FastAPI()


class RoleGroup(Enum):
    USER = ["USER", "PREMIUM_USER"]
    ADMIN = ["ADMIN"]


@app.get("/test")
def test(
    user: UserModel = Depends(RoleChecker(RoleGroup.USER.value)),
):
    return {"message": f"Hello, {user.username}!"}


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)

 

 

기본적으로 Depends()안에 들어가는 함수들은 동적으로 데이터(parameter)를 받지 못한다. 따라서, 이 코드의 핵심은 RoleChecker class를 이용해서, parameter를 사용할 수 있게 했다.

결과적으로, /test api는 사용자가 USER에 속하는 role을 가지고 있다면, API 사용을 허락하는 명시적이고, 유연성 있는 코드로 작성 되었다.

 

 

references

https://fastapi.tiangolo.com/advanced/advanced-dependencies/