utils/db.py - features/sqlalchemy - burevol/GymAssistant

utils/db.py
85 строк3 KB
Alexander Maximov
Добавлена возможность работы с PostgreSQL
09 май 2025, 00:02
09 май 2025, 00:02c9ea399
import os
from typing import Optional

from sqlalchemy import String, exists, select
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column

from utils.models import UserInDB


class Base(DeclarativeBase):
    pass


class User(Base):
    __tablename__ = "users"

    id: Mapped[int] = mapped_column(primary_key=True)
    username: Mapped[str] = mapped_column(String(50))
    hashed_password: Mapped[str] = mapped_column(String(60))
    disabled: Mapped[bool] = mapped_column()

    def __repr__(self):
        return f"<User {self.username}>"


class Db:
    def __init__(self):

        self.driver = os.getenv("DB_DRIVER", "sqlite")
        self.user = os.getenv("DB_USER", "test")
        self.password = os.getenv("DB_PASSWORD", "test")
        self.host = os.getenv("DB_HOST", "localhost")
        self.port = os.getenv("DB_PORT", "5432")
        self.database = os.getenv("DB_DATABASE", "mydatabase")

        self.url = f'{self.driver}://{self.user}:{self.password}@{self.host}:{self.port}/{self.database}'

        self.async_engine = create_async_engine(
            self.url,
            echo=True,  # Логирование SQL-запросов (для разработки)
            pool_size=5,  # Размер пула соединений
            max_overflow=10  # Максимальное количество соединений сверх pool_size
        )

        self.AsyncSessionLocal = async_sessionmaker(
            bind=self.async_engine,
            class_=AsyncSession,
            expire_on_commit=False,
            autocommit=False,
            autoflush=False
        )

    async def user_exists(self, username: str) -> bool:
        async with self.AsyncSessionLocal() as session:
            stmt = select(exists().where(User.username == username))
            result = await session.scalar(stmt)
        return result

    async def create_first_user(self):
        user_exists = await self.user_exists(username="johndoe")
        if not user_exists:
            async with self.AsyncSessionLocal() as session:
                # noinspection SpellCheckingInspection
                fake_user = User(username="johndoe",
                                 hashed_password="$2b$12$EixZaYVK1fsbw1ZfbX3OXePaWxn96p36WQoeG6Lruj3vjPGga31lW",
                                 disabled=False)
                session.add(fake_user)
                await session.commit()
                await session.refresh(fake_user)

    async def get_user(self, username: str) -> Optional[UserInDB]:
        user_exists = await self.user_exists(username=username)
        if user_exists:
            async with self.AsyncSessionLocal() as session:
                stmt = select(User).where(User.username == username)
                result = await session.scalar(stmt)
            return UserInDB.model_validate(result)

    async def delete_user(self, username: str):
        async with self.AsyncSessionLocal() as session:
            stmt = select(User).where(User.username == username)
            result = await session.scalar(stmt)
            await session.delete(result)
            await session.commit()