from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine

from config import cfg

engine = create_async_engine(cfg.db_url, echo=True)
async_session = async_sessionmaker(bind=engine, expire_on_commit=False)


class UnitOfWork:
    def __init__(self, session: AsyncSession) -> None:
        self.session = session

    async def commit(self) -> None:
        await self.session.commit()

    async def rollback(self) -> None:
        await self.session.rollback()

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, *_):
        if exc_type:
            await self.rollback()


async def get_db():
    async with async_session() as session:
        yield session


async def get_uow():
    async with async_session() as session, UnitOfWork(session) as uow:
        yield uow