<reponame>tropicoo/yt-dlp-bot
import re
import uuid
from typing import AsyncGenerator

import sqlalchemy as sa
from sqlalchemy import MetaData
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
from sqlalchemy.orm import (
    declarative_base,
    declared_attr,
    sessionmaker,
)
from sqlalchemy_utils import UUIDType

from yt_shared.config import (
    APPLICATION_NAME,
    SQLALCHEMY_DATABASE_URI_ASYNC,
    SQLALCHEMY_ECHO,
    SQLALCHEMY_EXPIRE_ON_COMMIT,
)

engine = create_async_engine(
    SQLALCHEMY_DATABASE_URI_ASYNC,
    echo=SQLALCHEMY_ECHO,
    pool_pre_ping=True,
    connect_args={'server_settings': {'application_name': APPLICATION_NAME}},
)

metadata = MetaData()
metadata.bind = engine
AsyncSessionLocal = sessionmaker(
    engine,
    class_=AsyncSession,
    expire_on_commit=SQLALCHEMY_EXPIRE_ON_COMMIT,
)


async def get_db() -> AsyncGenerator[AsyncSession, None]:
    async with AsyncSessionLocal() as session:
        yield session


class CustomBase:
    id: uuid.UUID = sa.Column(UUIDType(binary=False), primary_key=True,
                              default=uuid.uuid4)
    __name__: str

    @declared_attr
    def __tablename__(cls):  # noqa
        return re.sub(r'(?<!^)(?=[A-Z])', '_', cls.__name__).lower()


Base = declarative_base(cls=CustomBase, metadata=metadata)
