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