Spaces:
Sleeping
Sleeping
File size: 2,021 Bytes
39e8e23 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | """
S.T.E.W Database — SQLAlchemy async engine + session factory.
Supports PostgreSQL (production) and SQLite (dev/testing).
"""
import os
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker
from sqlalchemy.orm import DeclarativeBase
from sqlalchemy.pool import NullPool, AsyncAdaptedQueuePool
from server.config import get_settings
settings = get_settings()
def _get_async_url(url: str) -> str:
if url.startswith("postgresql://"):
return url.replace("postgresql://", "postgresql+asyncpg://", 1)
if url.startswith("postgres://"):
return url.replace("postgres://", "postgresql+asyncpg://", 1)
if url.startswith("sqlite:///") or url.startswith("sqlite+aiosqlite"):
if "sqlite+aiosqlite" not in url:
return url.replace("sqlite:///", "sqlite+aiosqlite:///", 1)
return url
return url
ASYNC_DATABASE_URL = _get_async_url(settings.DATABASE_URL)
IS_SQLITE = "sqlite" in ASYNC_DATABASE_URL
# SQLite doesn't support pool_size/max_overflow
if IS_SQLITE:
engine = create_async_engine(
ASYNC_DATABASE_URL,
echo=settings.DEBUG,
poolclass=NullPool,
)
else:
engine = create_async_engine(
ASYNC_DATABASE_URL,
echo=settings.DEBUG,
pool_pre_ping=True,
pool_size=5,
max_overflow=10,
)
AsyncSessionLocal = async_sessionmaker(
bind=engine,
class_=AsyncSession,
expire_on_commit=False,
autoflush=False,
autocommit=False,
)
class Base(DeclarativeBase):
pass
async def get_db() -> AsyncSession:
async with AsyncSessionLocal() as session:
try:
yield session
await session.commit()
except Exception:
await session.rollback()
raise
finally:
await session.close()
async def init_db():
"""Create all tables (used in dev/test; prod uses Alembic)."""
async with engine.begin() as conn:
await conn.run_sync(Base.metadata.create_all)
|