diff --git a/pilot/user/__init__.py b/pilot/user/__init__.py index 79cef7567..c45c75c85 100644 --- a/pilot/user/__init__.py +++ b/pilot/user/__init__.py @@ -3,7 +3,8 @@ from pilot.user.user_db import ( UserDao ) -from pilot.user.user_request import UserRequest, get_user_from_headers +from pilot.user.user_request import UserRequest +from pilot.user.auth import get_user_from_headers __all__ = [ "UserEntity", diff --git a/pilot/user/auth.py b/pilot/user/auth.py new file mode 100644 index 000000000..7652f2592 --- /dev/null +++ b/pilot/user/auth.py @@ -0,0 +1,15 @@ +from fastapi import Header +from pilot.user import UserDao, UserRequest + + +def get_user_from_headers(user_id: str = Header(...)): + try: + if user_id: + # 确保数据库存在当前用户,如果不存在则抛出异常 + user_dao = UserDao() + user_dao.get_by_user_id(user_id=user_id) + return UserRequest(user_id=user_id) + else: + raise f"User Not Login" + except Exception as e: + raise f"Authentication failed. {str(e)}" diff --git a/pilot/user/user_request.py b/pilot/user/user_request.py index 870fcfe55..180e760d8 100644 --- a/pilot/user/user_request.py +++ b/pilot/user/user_request.py @@ -1,6 +1,4 @@ -import json -from fastapi import Header from pydantic import BaseModel @@ -13,10 +11,3 @@ class UserRequest(BaseModel): nick_name: str = None email: str = None avatar_url: str = None - - -def get_user_from_headers(user_id: str = Header(...)): - try: - return UserRequest(user_id=user_id) - except Exception as e: - raise f"Authentication failed."