34 lines
879 B
Python
34 lines
879 B
Python
|
from collections.abc import AsyncIterator
|
||
|
|
|
||
|
|
from sqlalchemy.ext.asyncio import (
|
||
|
|
AsyncSession,
|
||
|
|
async_sessionmaker,
|
||
|
|
create_async_engine,
|
||
|
|
)
|
||
|
|
|
||
|
|
from ..config import settings
|
||
|
|
|
||
|
|
|
||
|
|
def _normalize_url(url: str) -> str:
|
||
|
|
"""Convert any postgres URL flavor to the asyncpg driver form."""
|
||
|
|
if url.startswith("postgres://"):
|
||
|
|
url = "postgresql://" + url[len("postgres://") :]
|
||
|
|
if url.startswith("postgresql://") and "+asyncpg" not in url:
|
||
|
|
url = "postgresql+asyncpg://" + url[len("postgresql://") :]
|
||
|
|
return url
|
||
|
|
|
||
|
|
|
||
|
|
engine = create_async_engine(
|
||
|
|
_normalize_url(settings.DATABASE_URL),
|
||
|
|
pool_size=10,
|
||
|
|
max_overflow=20,
|
||
|
|
pool_pre_ping=True,
|
||
|
|
)
|
||
|
|
|
||
|
|
SessionLocal = async_sessionmaker(engine, expire_on_commit=False, class_=AsyncSession)
|
||
|
|
|
||
|
|
|
||
|
|
async def get_db() -> AsyncIterator[AsyncSession]:
|
||
|
|
async with SessionLocal() as session:
|
||
|
|
yield session
|