Spaces:
Paused
Paused
| 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: | |
| 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) | |