blocking-ip/core/models/db_helper.py

58 lines
1.4 KiB
Python

from typing import AsyncGenerator
from sqlalchemy.ext.asyncio import (
AsyncSession,
create_async_engine,
async_sessionmaker,
AsyncEngine,
)
from core.config import settings
class DatabaseHelper:
def __init__(
self,
url: str,
echo: bool,
echo_pool: bool,
pool_size: int,
max_overflow: int,
) -> None:
self.engine: AsyncEngine = create_async_engine(
url=url,
echo=echo,
echo_pool=echo_pool,
pool_size=pool_size,
max_overflow=max_overflow,
)
self.session_factory: async_sessionmaker[AsyncSession] = async_sessionmaker(
bind=self.engine,
autoflush=False,
autocommit=False,
expire_on_commit=False,
)
# def get_scoped_session(self):
# session = async_scoped_session(
# session_factory=self.session_factory,
# scopefunc=current_task,
# )
# return session
async def dispose(self) -> None:
await self.engine.dispose()
async def session_getter(self) -> AsyncGenerator[AsyncSession, None]:
async with self.session_factory() as session:
yield session
db_helper = DatabaseHelper(
url=str(settings.db.url),
echo=settings.db.echo,
echo_pool=settings.db.echo_pool,
pool_size=settings.db.pool_size,
max_overflow=settings.db.max_overflow,
)