src/utils/migrate.py

"""
Run ALTER TABLE ADD COLUMN IF NOT EXISTS for columns added after initial deploy.
Safe to run on every startup — IF NOT EXISTS is a no-op when column already exists.
"""

from loguru import logger

_VECTOR_MIGRATIONS = [
    "CREATE EXTENSION IF NOT EXISTS vector",
    # vector columns already exist in Supabase — ALTER TABLE causes statement timeout
]

_VECTOR_INDEXES = [
    "CREATE INDEX IF NOT EXISTS ix_posts_content_embedding ON posts USING ivfflat (content_embedding vector_cosine_ops) WITH (lists = 10)",
    "CREATE INDEX IF NOT EXISTS ix_trends_topic_embedding ON trends USING ivfflat (topic_embedding vector_cosine_ops) WITH (lists = 10)",
]

_MIGRATIONS = [
    # posts — columns added in post-MVP schema evolution
    "ALTER TABLE posts ADD COLUMN IF NOT EXISTS post_type VARCHAR DEFAULT 'image'",
    "ALTER TABLE posts ADD COLUMN IF NOT EXISTS carousel_urls JSON",
    "ALTER TABLE posts ADD COLUMN IF NOT EXISTS video_url VARCHAR",
    "ALTER TABLE posts ADD COLUMN IF NOT EXISTS link_url VARCHAR",
    "ALTER TABLE posts ADD COLUMN IF NOT EXISTS link_title VARCHAR",
    "ALTER TABLE posts ADD COLUMN IF NOT EXISTS link_description VARCHAR",
    "ALTER TABLE posts ADD COLUMN IF NOT EXISTS approval_status VARCHAR DEFAULT 'pending'",
    "ALTER TABLE posts ADD COLUMN IF NOT EXISTS approval_notes TEXT",
    "ALTER TABLE posts ADD COLUMN IF NOT EXISTS approved_by VARCHAR",
    "ALTER TABLE posts ADD COLUMN IF NOT EXISTS approved_at TIMESTAMP",
    "ALTER TABLE posts ADD COLUMN IF NOT EXISTS token_count INTEGER",
    "ALTER TABLE posts ADD COLUMN IF NOT EXISTS generation_cost FLOAT",
    "ALTER TABLE posts ADD COLUMN IF NOT EXISTS hallucination_score FLOAT",
    "ALTER TABLE posts ADD COLUMN IF NOT EXISTS updated_at TIMESTAMP",
    # posts — generation_error for failed generating jobs
    "ALTER TABLE posts ADD COLUMN IF NOT EXISTS generation_error TEXT",
    # agent_versions — created_by added later
    "ALTER TABLE agent_versions ADD COLUMN IF NOT EXISTS created_by VARCHAR",
    # seed default agent and agent_version so posts can reference agent_version_id=1
    "INSERT INTO agents (id, name, is_active) SELECT 1, 'default', true WHERE NOT EXISTS (SELECT 1 FROM agents WHERE id = 1)",
    "INSERT INTO agent_versions (id, agent_id, version, config) SELECT 1, 1, 1, '{}'::json WHERE NOT EXISTS (SELECT 1 FROM agent_versions WHERE id = 1)",
    # workflow_executions — execution_metadata added later
    "ALTER TABLE workflow_executions ADD COLUMN IF NOT EXISTS execution_metadata JSON",
    """
    CREATE TABLE IF NOT EXISTS brand_voices (
        id SERIAL PRIMARY KEY,
        brand_name VARCHAR DEFAULT 'Orbis',
        voice TEXT,
        audience TEXT,
        tone_rules TEXT,
        banned_words TEXT,
        cta_rules TEXT,
        hashtag_rules TEXT,
        example_posts TEXT,
        is_active BOOLEAN DEFAULT TRUE,
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
    )
    """,
    "CREATE INDEX IF NOT EXISTS ix_brand_voice_active ON brand_voices(is_active)",
    # Performance indexes for common query patterns
    "CREATE INDEX IF NOT EXISTS ix_posts_trend_status ON posts(trend_id, status)",
    "CREATE INDEX IF NOT EXISTS ix_posts_scheduled ON posts(scheduled_at) WHERE scheduled_at IS NOT NULL",
    "CREATE INDEX IF NOT EXISTS ix_trends_status_score ON trends(status, score DESC)",
    "ALTER TABLE posts ADD COLUMN IF NOT EXISTS scheduled_at TIMESTAMP",
    # music tracks library
    """
    CREATE TABLE IF NOT EXISTS music_tracks (
        id SERIAL PRIMARY KEY,
        title VARCHAR NOT NULL,
        artist VARCHAR,
        genre VARCHAR,
        mood VARCHAR,
        bpm INTEGER,
        duration_secs INTEGER,
        preview_url VARCHAR,
        cover_url VARCHAR,
        tags JSON DEFAULT '[]',
        is_active BOOLEAN DEFAULT TRUE,
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
    )
    """,
    "ALTER TABLE posts ADD COLUMN IF NOT EXISTS music_id INTEGER REFERENCES music_tracks(id) ON DELETE SET NULL",
    # agent_findings — AI agent team daily analysis results
    """
    CREATE TABLE IF NOT EXISTS agent_findings (
        id SERIAL PRIMARY KEY,
        agent_role VARCHAR NOT NULL,
        category VARCHAR,
        severity VARCHAR DEFAULT 'info',
        title VARCHAR NOT NULL,
        body TEXT,
        action_items JSON DEFAULT '[]',
        data_snapshot JSON DEFAULT '{}',
        run_date DATE NOT NULL,
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
    )
    """,
    "CREATE INDEX IF NOT EXISTS ix_agent_findings_run_date ON agent_findings(run_date)",
    "CREATE INDEX IF NOT EXISTS ix_agent_findings_agent_role ON agent_findings(agent_role)",
    "CREATE INDEX IF NOT EXISTS ix_agent_findings_severity ON agent_findings(severity)",
    # post quality scoring
    "ALTER TABLE posts ADD COLUMN IF NOT EXISTS quality_score INTEGER",
    "ALTER TABLE posts ADD COLUMN IF NOT EXISTS quality_breakdown JSON",
    # brand knowledge base
    """
    CREATE TABLE IF NOT EXISTS brand_knowledge (
        id SERIAL PRIMARY KEY,
        title VARCHAR NOT NULL,
        content TEXT NOT NULL,
        content_type VARCHAR DEFAULT 'general',
        embedding vector(1536),
        is_active BOOLEAN DEFAULT TRUE,
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
    )
    """,
    "CREATE INDEX IF NOT EXISTS ix_brand_knowledge_active ON brand_knowledge(is_active)",
    "CREATE INDEX IF NOT EXISTS ix_brand_knowledge_type ON brand_knowledge(content_type)",
    # provider_keys — encrypted LLM API keys manageable from Settings UI
    """
    CREATE TABLE IF NOT EXISTS provider_keys (
        id SERIAL PRIMARY KEY,
        provider VARCHAR NOT NULL UNIQUE,
        display_name VARCHAR,
        api_key_enc TEXT,
        api_base VARCHAR,
        is_enabled BOOLEAN DEFAULT TRUE,
        models JSON DEFAULT '[]',
        updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
    )
    """,
    "ALTER TABLE provider_keys ADD COLUMN IF NOT EXISTS models JSON DEFAULT '[]'",
    "ALTER TABLE provider_keys ADD COLUMN IF NOT EXISTS disabled_models JSON DEFAULT '[]'",
    # brand_voices — auto hashtags + caption templates
    "ALTER TABLE brand_voices ADD COLUMN IF NOT EXISTS auto_hashtags TEXT",
    "ALTER TABLE brand_voices ADD COLUMN IF NOT EXISTS caption_templates JSON DEFAULT '[]'",
    # posts — monthly budget tracking + engagement metrics
    "ALTER TABLE posts ADD COLUMN IF NOT EXISTS budget_month VARCHAR(7)",
    "ALTER TABLE posts ADD COLUMN IF NOT EXISTS ig_likes INTEGER",
    "ALTER TABLE posts ADD COLUMN IF NOT EXISTS ig_comments INTEGER",
    "ALTER TABLE posts ADD COLUMN IF NOT EXISTS ig_reach INTEGER",
    "ALTER TABLE posts ADD COLUMN IF NOT EXISTS ig_saves INTEGER",
    "ALTER TABLE posts ADD COLUMN IF NOT EXISTS ig_shares INTEGER",
    "ALTER TABLE posts ADD COLUMN IF NOT EXISTS engagement_fetched_at TIMESTAMP",
]

_MUSIC_SEED = [
    ("Upbeat Rush", "Studio Orbit", "electronic", "energetic", 128, 180, '["workout","fitness","sport"]'),
    ("Power Drive", "Studio Orbit", "pop", "energetic", 135, 165, '["energy","hustle","motivation"]'),
    ("Momentum", "Studio Orbit", "electronic", "energetic", 130, 200, '["gym","run","power"]'),
    ("Morning Breath", "Studio Orbit", "ambient", "calm", 70, 210, '["meditation","yoga","peace"]'),
    ("Soft Focus", "Studio Orbit", "acoustic", "calm", 75, 195, '["relax","nature","serene"]'),
    ("Gentle Flow", "Studio Orbit", "ambient", "calm", 65, 225, '["wellness","mindful","breathe"]'),
    ("Summer Haze", "Studio Orbit", "pop", "upbeat", 112, 175, '["lifestyle","fashion","fun"]'),
    ("Golden Vibes", "Studio Orbit", "pop", "upbeat", 118, 185, '["trending","happy","style"]'),
    ("Fresh Start", "Studio Orbit", "pop", "upbeat", 110, 170, '["beauty","positive","joy"]'),
    ("Rise Above", "Studio Orbit", "cinematic", "inspirational", 90, 240, '["success","growth","goal"]'),
    ("New Chapter", "Studio Orbit", "cinematic", "inspirational", 85, 220, '["journey","dream","achieve"]'),
    ("Limitless", "Studio Orbit", "cinematic", "inspirational", 88, 235, '["startup","motivate","rise"]'),
    ("Digital Pulse", "Studio Orbit", "electronic", "modern", 130, 190, '["tech","ai","digital"]'),
    ("Future State", "Studio Orbit", "electronic", "modern", 125, 195, '["innovation","code","data"]'),
    ("Kitchen Stories", "Studio Orbit", "acoustic", "warm", 80, 185, '["food","cook","recipe"]'),
    ("Cozy Corner", "Studio Orbit", "acoustic", "warm", 72, 200, '["home","family","comfort"]'),
    ("Golden Hour", "Studio Orbit", "acoustic", "romantic", 72, 215, '["love","couple","romance"]'),
    ("Sweet Embrace", "Studio Orbit", "acoustic", "romantic", 68, 205, '["heart","wedding","together"]'),
    ("Smooth Edge", "Studio Orbit", "jazz", "sophisticated", 95, 195, '["luxury","premium","business"]'),
    ("Bounce", "Studio Orbit", "pop", "playful", 116, 160, '["funny","fun","play","kids"]'),
]


async def run_migrations():
    import time
    from src.utils.database import engine
    from sqlalchemy import text

    # Hard time budget — migrations are mostly no-ops on an established DB,
    # so if they hang past this we just move on and the server still starts.
    BUDGET_SECS = 60
    start = time.monotonic()

    for stmt in _VECTOR_MIGRATIONS:
        if time.monotonic() - start > BUDGET_SECS:
            logger.warning("[Migrate] Vector budget exceeded — skipping rest")
            break
        try:
            async with engine.begin() as conn:
                await conn.execute(text(stmt))
        except Exception as e:
            logger.warning(f"[Migrate] Vector stmt skipped: {e!s:.120}")
    logger.info("[Migrate] Vector extension and columns applied")

    for stmt in _MIGRATIONS:
        if time.monotonic() - start > BUDGET_SECS:
            logger.warning("[Migrate] Migration budget exceeded — assuming schema is current")
            break
        try:
            async with engine.begin() as conn:
                await conn.execute(text(stmt))
        except Exception as e:
            logger.warning(f"[Migrate] Stmt skipped (non-fatal): {e!s:.120}")
    logger.info("[Migrate] Column migrations applied")

    # ivfflat indexes require at least 1 row — best-effort
    try:
        async with engine.begin() as conn:
            for stmt in _VECTOR_INDEXES:
                await conn.execute(text(stmt))
        logger.info("[Migrate] Vector indexes applied")
    except Exception as e:
        logger.warning(f"[Migrate] Vector index creation skipped (table may be empty): {e}")

    await _seed_music_tracks()


async def _seed_music_tracks():
    from src.utils.database import engine
    from sqlalchemy import text
    try:
        async with engine.begin() as conn:
            count = (await conn.execute(text("SELECT COUNT(*) FROM music_tracks"))).scalar()
            if count == 0:
                for (title, artist, genre, mood, bpm, dur, tags) in _MUSIC_SEED:
                    await conn.execute(text(
                        "INSERT INTO music_tracks (title, artist, genre, mood, bpm, duration_secs, tags) "
                        "VALUES (:t, :a, :g, :m, :b, :d, CAST(:tg AS json))"
                    ), {"t": title, "a": artist, "g": genre, "m": mood, "b": bpm, "d": dur, "tg": tags})
                logger.info(f"[Migrate] Seeded {len(_MUSIC_SEED)} music tracks")
    except Exception as e:
        logger.warning(f"[Migrate] Music seed skipped: {e}")