You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
46 lines
1.5 KiB
46 lines
1.5 KiB
from sqlalchemy.ext.asyncio import create_async_engine, AsyncEngine, async_sessionmaker, AsyncSession
|
|
from typing import AsyncGenerator
|
|
from core.settings import settings
|
|
class DatabaseHelper:
|
|
def __init__(
|
|
self,
|
|
url: str,
|
|
echo:bool=False,
|
|
echo_pool:bool=False,
|
|
pool_size: int = 5,
|
|
max_overflow: int =10,
|
|
) -> None:
|
|
try:
|
|
self.engine: AsyncEngine = create_async_engine(
|
|
url=url,
|
|
echo=echo,
|
|
echo_pool=echo_pool,
|
|
pool_size=pool_size,
|
|
max_overflow = max_overflow
|
|
)
|
|
except TypeError:
|
|
# The pool_size argument won't work for the default SQLite setup in SQLAlchemy 0.7, try without
|
|
self.engine: AsyncEngine = create_async_engine(url=url)
|
|
|
|
self.session_factory: async_sessionmaker[AsyncSession] = async_sessionmaker(
|
|
bind= self.engine,
|
|
autoflush=False,
|
|
autocommit=False,
|
|
expire_on_commit=False
|
|
)
|
|
|
|
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 = settings.database.url,
|
|
echo = settings.database.echo,
|
|
echo_pool = settings.database.echo_pool,
|
|
pool_size = settings.database.pool_size,
|
|
max_overflow = settings.database.max_overflow,
|
|
|
|
)
|