39 lines
907 B
Python
39 lines
907 B
Python
import asyncio
|
|
from contextlib import asynccontextmanager
|
|
from typing import AsyncGenerator, AsyncIterator
|
|
|
|
from alembic import command
|
|
from alembic.config import Config
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from src.core.config import settings
|
|
from src.db.base import SessionLocal
|
|
|
|
|
|
@asynccontextmanager
|
|
async def db_session() -> AsyncGenerator[AsyncSession, None]:
|
|
db = SessionLocal()
|
|
try:
|
|
yield db
|
|
await db.commit()
|
|
except Exception:
|
|
await db.rollback()
|
|
raise
|
|
finally:
|
|
await db.close()
|
|
|
|
|
|
async def get_db() -> AsyncGenerator[AsyncSession, None]:
|
|
async with db_session() as db:
|
|
yield db
|
|
|
|
|
|
async def run_migrations() -> None:
|
|
alembic_cfg = Config("alembic.ini")
|
|
await asyncio.to_thread(command.upgrade, alembic_cfg, "head")
|
|
|
|
|
|
async def create_tables() -> None:
|
|
if not settings.DEBUG:
|
|
await run_migrations()
|