from collections.abc import AsyncGenerator from pathlib import Path from sqlalchemy import event from sqlalchemy.ext.asyncio import ( AsyncSession, async_sessionmaker, create_async_engine, ) from sqlalchemy.orm import DeclarativeBase from app.config import get_settings settings = get_settings() _is_sqlite = settings.database_url.startswith("sqlite") if _is_sqlite: db_path = settings.database_url.split("///")[-1] Path(db_path).parent.mkdir(parents=True, exist_ok=True) engine = create_async_engine( settings.database_url, echo=False, future=True, connect_args={"timeout": 30} if _is_sqlite else {}, ) if _is_sqlite: @event.listens_for(engine.sync_engine, "connect") def _sqlite_pragmas(dbapi_connection, connection_record) -> None: cursor = dbapi_connection.cursor() cursor.execute("PRAGMA busy_timeout=30000") cursor.execute("PRAGMA foreign_keys=ON") cursor.close() AsyncSessionLocal = async_sessionmaker( bind=engine, class_=AsyncSession, expire_on_commit=False, autoflush=False, ) class Base(DeclarativeBase): pass async def get_db() -> AsyncGenerator[AsyncSession, None]: async with AsyncSessionLocal() as session: yield session async def init_db() -> None: from app import models async with engine.begin() as conn: await conn.run_sync(Base.metadata.create_all)