# Copyright 2026 zhaoxi826 # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from pretor.core.database.table.user import User from sqlmodel import select from pretor.utils.error import UserNotExistError, UserPasswordError from pretor.core.database.database_exception import database_exception from pretor.core.database.table.user import UserAuthority from pretor.utils.access import Accessor class AuthDatabase: def __init__(self, async_session_maker): self.async_session_maker = async_session_maker @database_exception async def add_user(self, user_name: str, hashed_password: str) -> User: from ulid import ULID async with self.async_session_maker() as session: # Check if any users exist statement = select(User).limit(1) results = await session.execute(statement) existing_user = results.first() authority = UserAuthority.USER if existing_user is None: authority = UserAuthority.SUPER_ADMINISTRATOR user = User( user_id=str(ULID()), user_name=user_name, hashed_password=hashed_password, user_authority=authority ) session.add(user) await session.commit() await session.refresh(user) return user @database_exception async def change_password(self, user_name, old_password, new_password) -> User: async with self.async_session_maker() as session: statement = select(User).where(User.user_name == user_name) results = await session.execute(statement) user = results.scalar_one_or_none() if user is None: raise UserNotExistError() if not Accessor.verify_password(old_password, user.hashed_password): raise UserPasswordError() user.hashed_password = new_password session.add(user) await session.commit() await session.refresh(user) return user @database_exception async def delete_user(self, user_name: str) -> None: async with self.async_session_maker() as session: statement = select(User).where(User.user_name == user_name) results = await session.execute(statement) user = results.scalar_one_or_none() if user is None: raise UserNotExistError() session.delete(user) await session.commit() @database_exception async def delete_user_by_id(self, user_id: str) -> None: async with self.async_session_maker() as session: user = await session.get(User, user_id) if user is None: raise UserNotExistError() session.delete(user) await session.commit() @database_exception async def login_user(self, user_name: str) -> str: async with self.async_session_maker() as session: statement = select(User).where(User.user_name == user_name) results = await session.execute(statement) user = results.scalar_one_or_none() if user is None: raise UserNotExistError() return user @database_exception async def get_all_users(self) -> list[User]: async with self.async_session_maker() as session: statement = select(User) results = await session.execute(statement) users = results.scalars().all() return list(users) @database_exception async def get_user_authority(self, user_id: str) -> UserAuthority: async with self.async_session_maker() as session: user = await session.get(User, user_id) if user is None: raise UserNotExistError() return user.user_authority @database_exception async def change_user_authority(self, user_id: str, new_authority: UserAuthority) -> User: """ Changes the authority level of a specific user. Args: user_id: The ID of the user whose authority is to be changed. new_authority: The new authority level to assign to the user. Returns: User: The updated user object. Raises: UserNotExistError: If the specified user does not exist. """ async with self.async_session_maker() as session: user = await session.get(User, user_id) if user is None: raise UserNotExistError() user.user_authority = new_authority session.add(user) await session.commit() await session.refresh(user) return user