ai-rpg/backend/app/migrations/init_db.py

"""Database initialization: create all tables and seed defaults.

Idempotent: safe to call from multiple processes (backend lifespan + worker
startup) thanks to a PostgreSQL advisory lock that serializes the seeding
phase. `create_all` itself is already `CREATE TABLE IF NOT EXISTS`, so the
only race is on seed inserts — guarded by `pg_advisory_xact_lock` plus
per-row IntegrityError handling.
"""
from __future__ import annotations

import asyncio
import json
from pathlib import Path

from sqlalchemy import select, text
from sqlalchemy.exc import IntegrityError

from app.db import AsyncSessionLocal, Base, engine
from app.models import GlossaryEntry, Preset, Setting, User
from app.config import settings
from app.logging_setup import get_logger, setup_logging
from app.core.security import hash_password
from app.prompts.fantasy_preset import FANTASY_PRESET_RU, FANTASY_PRESET_EN

log = get_logger("migrations")


# Stable advisory lock key so backend + worker don't race on seeding.
# (key1, key2) — arbitrary 64-bit integers, kept constant across runs.
_ADVISORY_LOCK_KEY = (42424201, 1)


DEFAULT_SETTINGS = [
    ("llm.base_url", settings.default_llm_base_url, "OpenAI-compatible base URL"),
    ("llm.api_key", settings.default_llm_api_key, "API key for LLM endpoint"),
    ("llm.model", settings.default_llm_model, "Default model name"),
    ("llm.temperature", 0.7, "Temperature for orchestrator"),
    ("llm.step_temperature", 0.85, "Temperature for narrative step writer"),
    ("llm.summary_temperature", 0.3, "Temperature for summarizer"),
    ("llm.max_tokens", 1024, "Max tokens per LLM response"),
    ("llm.request_timeout", 120, "LLM request timeout, seconds"),
    ("llm.streaming", True, "Whether to use streaming responses"),
    ("context.recent_messages", settings.default_recent_messages, "Guaranteed recent messages in prompt"),
    ("context.compress_threshold", settings.default_compress_threshold, "Trigger compression at this count"),
    ("context.summary_messages", settings.default_summary_messages, "Number of messages per summary block"),
    ("context.max_tokens_total", 6000, "Soft token budget for context window (small models)"),
    ("triggers.enabled", True, "Enable deferred trigger processing"),
    ("triggers.check_interval", 30, "Trigger checker interval, seconds"),
    # Embeddings / RAG
    ("embedding.provider", settings.default_embedding_provider, "Embeddings provider: 'hash' (offline fallback) or 'openai' (real semantic embeddings)"),
    ("embedding.base_url", settings.default_embedding_base_url, "OpenAI-compatible embeddings base URL. Empty = reuse llm.base_url"),
    ("embedding.api_key", settings.default_embedding_api_key, "API key for embeddings endpoint. Empty = reuse llm.api_key"),
    ("embedding.model", settings.default_embedding_model, "Embedding model name (e.g. text-embedding-3-small, bge-m3, nomic-embed-text)"),
    ("embedding.dim", settings.default_embedding_dim, "Vector dimension. 0 = auto-probe from endpoint on first use"),
    ("embedding.request_timeout", settings.default_embedding_request_timeout, "Embeddings request timeout, seconds"),
]


async def init_db() -> None:
    setup_logging()
    log.info("creating_tables")
    # CREATE TABLE IF NOT EXISTS — safe to run concurrently.
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
    log.info("tables_ready")

    # Seed phase: serialize across processes via PG advisory transaction lock.
    # On non-PG backends (SQLite for tests) the lock statement is a no-op
    # (we catch the error and proceed without locking).
    try:
        async with AsyncSessionLocal() as session:
            await session.execute(
                text("SELECT pg_advisory_xact_lock(:k1, :k2)").bindparams(
                    k1=_ADVISORY_LOCK_KEY[0], k2=_ADVISORY_LOCK_KEY[1]
                )
            )
            await _seed_settings(session)
            await _seed_builtin_presets(session)
            await session.commit()
    except Exception as e:
        # Non-PG backend (SQLite) or transient error — retry without the lock.
        log.warning("advisory_lock_unavailable_proceeding", error=f"{type(e).__name__}: {e}")
        async with AsyncSessionLocal() as session:
            await _seed_settings(session)
            await _seed_builtin_presets(session)
            await session.commit()

    # Ensure admin_setup_token is set; if empty, generate and print
    token = settings.admin_setup_token.strip()
    if not token:
        import secrets as _s
        token = _s.token_urlsafe(24)
        async with AsyncSessionLocal() as session:
            existing = await session.execute(select(Setting).where(Setting.key == "admin.setup_token"))
            existing_obj = existing.scalars().first()
            if existing_obj is None:
                session.add(Setting(key="admin.setup_token", value=token, description="One-time token for /admin/setup"))
                try:
                    await session.commit()
                    print("=" * 60)
                    print("ADMIN SETUP TOKEN (use at /admin/setup):")
                    print(token)
                    print("=" * 60)
                    log.info("admin_setup_token_generated")
                except IntegrityError:
                    # Another process inserted it concurrently — fine.
                    await session.rollback()
                    log.info("admin_setup_token_already_set")
            else:
                log.info("admin_setup_token_already_set")


async def _seed_settings(session) -> None:
    """Insert default settings that don't yet exist (per-row, race-safe)."""
    result = await session.execute(select(Setting).limit(1))
    if result.scalars().first() is not None:
        log.info("settings_already_exist")
        return
    seeded = 0
    for key, value, desc in DEFAULT_SETTINGS:
        # Check existence per-row to avoid IntegrityError on concurrent inserts
        existing = await session.execute(select(Setting).where(Setting.key == key))
        if existing.scalars().first() is not None:
            continue
        session.add(Setting(key=key, value=value, description=desc))
        seeded += 1
    if seeded:
        try:
            await session.commit()
            log.info("settings_seeded", count=seeded)
        except IntegrityError:
            await session.rollback()
            log.info("settings_seed_skipped_concurrent")
    else:
        log.info("settings_already_exist")


async def _seed_builtin_presets(session) -> None:
    """Insert built-in presets if none exist yet."""
    result = await session.execute(select(Preset).where(Preset.is_builtin.is_(True)))
    if result.scalars().first() is not None:
        log.info("builtin_presets_already_exist")
        return
    for preset_def in (FANTASY_PRESET_RU, FANTASY_PRESET_EN):
        # Check by slug to avoid race on unique constraint
        existing = await session.execute(select(Preset).where(Preset.slug == preset_def["slug"]))
        if existing.scalars().first() is not None:
            continue
        session.add(Preset(
            slug=preset_def["slug"],
            title=preset_def["title"],
            description=preset_def["description"],
            language=preset_def["language"],
            is_public=True,
            is_builtin=True,
            payload=preset_def["payload"],
        ))
    try:
        await session.commit()
        log.info("builtin_presets_seeded")
    except IntegrityError:
        await session.rollback()
        log.info("builtin_presets_seed_skipped_concurrent")


if __name__ == "__main__":
    asyncio.run(init_db())