File size: 999 Bytes
054900e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 |
from __future__ import annotations
from typing import TYPE_CHECKING
from uuid import uuid4
from asyncpg import Connection
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine
from bot.core.config import settings
if TYPE_CHECKING:
from sqlalchemy.engine.url import URL
class CConnection(Connection): # type: ignore
def _get_unique_id(self, prefix: str) -> str:
return f"__asyncpg_{prefix}_{uuid4()}__"
def get_engine(url: URL | str = settings.database_url) -> AsyncEngine:
return create_async_engine(
url=url,
echo=settings.DEBUG,
pool_size=0,
connect_args={
"connection_class": CConnection,
},
)
def get_sessionmaker(engine: AsyncEngine) -> async_sessionmaker[AsyncSession]:
return async_sessionmaker(bind=engine, autoflush=False, expire_on_commit=False)
db_url = settings.database_url
engine = get_engine(url=db_url)
sessionmaker = get_sessionmaker(engine)
|