Навёл немного красоты

chat_test/app/users/auth.py

@@ -5,9 +5,10 @@ from passlib.context import CryptContext
 from pydantic import EmailStr
 
 from app.config import settings
-from app.exceptions import UserDontHavePermissionException
+from app.exceptions import UserDontHavePermissionException, IncorrectAuthDataException, UserAlreadyExistsException
 from app.users.dao import UserDAO
 from app.users.models import Users
+from app.users.schemas import SUserRegister
 
 pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
 
@@ -51,6 +52,24 @@ async def authenticate_user_by_username(username: str, password: str) -> Users |
     return user
 
 
+async def authenticate_user(email_or_username: str, password: str) -> Users:
+    user = await authenticate_user_by_email(email_or_username, password)
+    if not user:
+        user = await authenticate_user_by_username(email_or_username, password)
+    if not user:
+        raise IncorrectAuthDataException
+    return user
+
+
+async def check_existing_user(user_data: SUserRegister) -> None:
+    existing_user = await UserDAO.find_one_or_none(email=user_data.email)
+    if existing_user:
+        raise UserAlreadyExistsException
+    existing_user = await UserDAO.find_one_or_none(username=user_data.username)
+    if existing_user:
+        raise UserAlreadyExistsException
+
+
 async def get_user_allowed_chats_id(user_id: int):
     user_allowed_chats = await UserDAO.get_user_allowed_chats(user_id)
     user_allowed_chats_id = (chat['chat_id'] for chat in user_allowed_chats)

chat_test/app/users/dao.py

@@ -1,8 +1,7 @@
 from datetime import datetime
 
-from sqlalchemy import update, select, insert, and_, func, literal, text
+from sqlalchemy import update, select, insert, and_, func, text
 from sqlalchemy.exc import SQLAlchemyError
-from sqlalchemy.dialects.postgresql import INTERVAL
 
 from app.dao.base import BaseDAO
 from app.database import async_session_maker, engine

chat_test/app/users/router.py

@@ -2,10 +2,11 @@ from fastapi import APIRouter, Response, Depends
 from fastapi.responses import RedirectResponse
 from starlette import status
 
-from app.exceptions import UserAlreadyExistsException, IncorrectAuthDataException, UsernameAlreadyInUseException, \
+from app.exceptions import UserAlreadyExistsException, UsernameAlreadyInUseException, \
     IncorrectPasswordException, PasswordsМismatchException, WrongCodeException
 from app.users.auth import get_password_hash, authenticate_user_by_email, authenticate_user_by_username, \
-    create_access_token, verify_password, REGISTRATED_USER, get_user_codes_list, VERIFICATED_USER
+    create_access_token, verify_password, REGISTRATED_USER, get_user_codes_list, VERIFICATED_USER, authenticate_user, \
+    check_existing_user
 from app.users.dao import UserDAO, UserCodesDAO
 from app.users.dependencies import get_current_user
 from app.users.models import Users
@@ -31,12 +32,7 @@ async def get_all_users():
 
 @router.post("/register", response_model=dict[str, str])
 async def register_user(response: Response, user_data: SUserRegister):
-    existing_user = await UserDAO.find_one_or_none(email=user_data.email)
+    await check_existing_user(user_data)
-    if existing_user:
-        raise UserAlreadyExistsException
-    existing_user = await UserDAO.find_one_or_none(username=user_data.username)
-    if existing_user:
-        raise UserAlreadyExistsException
     hashed_password = get_password_hash(user_data.password)
     user_id = await UserDAO.add(
         email=user_data.email,
@@ -69,11 +65,7 @@ async def email_verification(user_code: str, user: Users = Depends(get_current_u
 
 @router.post("/login", response_model=dict[str, str])
 async def login_user(response: Response, user_data: SUserLogin):
-    user = await authenticate_user_by_email(user_data.email_or_username, user_data.password)
+    user = await authenticate_user(user_data.email_or_username, user_data.password)
-    if not user:
-        user = await authenticate_user_by_username(user_data.email_or_username, user_data.password)
-    if not user:
-        raise IncorrectAuthDataException
     access_token = create_access_token({"sub": str(user.id)})
     response.set_cookie("black_phoenix_access_token", access_token, httponly=True)
     return {"access_token": access_token}