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