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 trigger firing on in-game time changes"),
    # 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"),
]


# Keys whose values come from environment variables (via Settings fields).
# These are re-applied on EVERY startup so .env is the source of truth.
# Admin-panel changes to these keys are runtime overrides that get reset on
# restart unless the operator also updates .env.
ENV_DERIVED_SETTING_KEYS = {
    "llm.base_url",
    "llm.api_key",
    "llm.model",
    "embedding.provider",
    "embedding.base_url",
    "embedding.api_key",
    "embedding.model",
    "embedding.dim",
    "embedding.request_timeout",
}


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 _sync_env_derived_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 _sync_env_derived_settings(session)
            await _seed_builtin_presets(session)
            await session.commit()

    # Ensure admin_setup_token is set and print it on every startup.
    #
    # The admin-setup endpoint refuses to create a second admin (see app/api/auth.py),
    # so it's safe to always print the token — even after an admin exists, the token
    # is useless. We print on every startup (not just first run) so the operator can
    # always find the URL in the logs without having to dig through old logs.
    token = settings.admin_setup_token.strip()
    if not token:
        # No token forced via env — generate one and persist it (idempotent).
        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()
                except IntegrityError:
                    # Another process inserted it concurrently — re-read.
                    await session.rollback()
                    await session.rollback()
                    existing = await session.execute(select(Setting).where(Setting.key == "admin.setup_token"))
                    existing_obj = existing.scalars().first()
                    if existing_obj is not None:
                        token = str(existing_obj.value)
            else:
                # Use the persisted token (env was empty, DB has one).
                token = str(existing_obj.value)
    # Always print — operator convenience.
    print("=" * 60)
    print("ADMIN SETUP URL:")
    print(f"  /admin/setup")
    print("ADMIN SETUP TOKEN:")
    print(f"  {token}")
    print("=" * 60)
    log.info("admin_setup_token_printed")


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 _sync_env_derived_settings(session) -> None:
    """Re-apply env-derived setting values from .env on every startup.

    This makes .env the source of truth for these keys: changing .env and
    restarting the container takes effect immediately. Admin-panel edits to
    these keys are runtime overrides that are reset on the next restart
    (unless the operator also updates .env).

    Only the env-derived keys (see ENV_DERIVED_SETTING_KEYS) are touched;
    other settings (temperature, context params, etc.) are preserved as
    configured via the admin panel.
    """
    env_values = {key: value for key, value, _desc in DEFAULT_SETTINGS if key in ENV_DERIVED_SETTING_KEYS}
    updated = 0
    for key, new_value in env_values.items():
        result = await session.execute(select(Setting).where(Setting.key == key))
        row = result.scalars().first()
        if row is None:
            # Shouldn't happen (seeded above) but handle defensively.
            session.add(Setting(key=key, value=new_value, description="Env-derived"))
            updated += 1
        else:
            if row.value != new_value:
                log.info(
                    "env_setting_resynced",
                    key=key,
                    old_value=str(row.value)[:80],
                    new_value=str(new_value)[:80],
                )
                row.value = new_value
                updated += 1
    if updated:
        try:
            await session.commit()
            log.info("env_settings_synced", count=updated)
        except IntegrityError:
            await session.rollback()
            log.warning("env_settings_sync_failed_concurrent")


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())