Imago/app/database.py

79 lines
2.4 KiB
Python

"""
Configuration SQLAlchemy — session async
"""
import logging
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
from sqlalchemy.orm import DeclarativeBase
from app.config import settings
logger = logging.getLogger(__name__)
engine = create_async_engine(
settings.DATABASE_URL,
echo=settings.DEBUG,
future=True,
)
AsyncSessionLocal = async_sessionmaker(
bind=engine,
class_=AsyncSession,
expire_on_commit=False,
autoflush=False,
autocommit=False,
)
class Base(DeclarativeBase):
pass
async def get_db() -> AsyncSession:
"""Dependency FastAPI — injecte une session DB dans chaque requête."""
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():
"""Crée toutes les tables et initialise un client par défaut si nécessaire."""
import secrets
import hashlib
from sqlalchemy import select
from app.models.client import APIClient, ClientPlan
async with engine.begin() as conn:
from app.models import image # noqa: F401
from app.models import client # noqa: F401
await conn.run_sync(Base.metadata.create_all)
# Vérifier s'il y a déjà des clients
async with AsyncSessionLocal() as session:
result = await session.execute(select(APIClient).limit(1))
if result.scalar_one_or_none() is None:
# Table vide -> Création du client bootstrap
raw_key = secrets.token_urlsafe(32)
key_hash = hashlib.sha256(raw_key.encode("utf-8")).hexdigest()
bootstrap_client = APIClient(
name="Default Admin",
api_key_hash=key_hash,
scopes=["images:read", "images:write", "images:delete", "ai:use", "admin"],
plan=ClientPlan.PREMIUM,
)
session.add(bootstrap_client)
await session.commit()
msg = f"Bootstrap client created! ID: {bootstrap_client.id} | API_KEY: {raw_key}"
logger.info(msg, extra={
"client_id": bootstrap_client.id,
"api_key": raw_key,
"warning": "Notez cette clé ! Elle ne sera plus affichée.",
})