pounce/backend/app/db_migrations.py

"""
Lightweight, idempotent DB migrations.

This project historically used `Base.metadata.create_all()` for bootstrapping new installs.
That does NOT handle schema evolution on existing databases. For performance-related changes
(indexes, new optional columns), we apply additive migrations on startup.

Important:
- Only additive changes (ADD COLUMN / CREATE INDEX) should live here.
- Operations must be idempotent (safe to run on every startup).
"""

from __future__ import annotations

import logging
from typing import Any

from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncConnection

logger = logging.getLogger(__name__)


async def _sqlite_table_exists(conn: AsyncConnection, table: str) -> bool:
    res = await conn.execute(
        text("SELECT 1 FROM sqlite_master WHERE type='table' AND name=:name LIMIT 1"),
        {"name": table},
    )
    return res.scalar() is not None


async def _sqlite_has_column(conn: AsyncConnection, table: str, column: str) -> bool:
    res = await conn.execute(text(f"PRAGMA table_info({table})"))
    rows = res.fetchall()
    # PRAGMA table_info: (cid, name, type, notnull, dflt_value, pk)
    return any(r[1] == column for r in rows)


async def _postgres_table_exists(conn: AsyncConnection, table: str) -> bool:
    # to_regclass returns NULL if the relation does not exist
    res = await conn.execute(text("SELECT to_regclass(:name)"), {"name": table})
    return res.scalar() is not None


async def _postgres_has_column(conn: AsyncConnection, table: str, column: str) -> bool:
    res = await conn.execute(
        text(
            """
            SELECT 1
            FROM information_schema.columns
            WHERE table_schema = current_schema()
              AND table_name = :table
              AND column_name = :column
            LIMIT 1
            """
        ),
        {"table": table, "column": column},
    )
    return res.scalar() is not None


async def _table_exists(conn: AsyncConnection, table: str) -> bool:
    dialect = conn.engine.dialect.name
    if dialect == "sqlite":
        return await _sqlite_table_exists(conn, table)
    return await _postgres_table_exists(conn, table)


async def _has_column(conn: AsyncConnection, table: str, column: str) -> bool:
    dialect = conn.engine.dialect.name
    if dialect == "sqlite":
        return await _sqlite_has_column(conn, table, column)
    return await _postgres_has_column(conn, table, column)


async def apply_migrations(conn: AsyncConnection) -> None:
    """
    Apply idempotent migrations.

    Called on startup after `create_all()` to keep existing DBs up-to-date.
    """
    dialect = conn.engine.dialect.name
    logger.info("DB migrations: starting (dialect=%s)", dialect)

    # ------------------------------------------------------------------
    # 1) domain_auctions.pounce_score (enables DB-level sorting/pagination)
    # ------------------------------------------------------------------
    if await _table_exists(conn, "domain_auctions"):
        if not await _has_column(conn, "domain_auctions", "pounce_score"):
            logger.info("DB migrations: adding column domain_auctions.pounce_score")
            await conn.execute(text("ALTER TABLE domain_auctions ADD COLUMN pounce_score INTEGER"))
        # Index for feed ordering
        await conn.execute(
            text("CREATE INDEX IF NOT EXISTS ix_domain_auctions_pounce_score ON domain_auctions(pounce_score)")
        )

    # ---------------------------------------------------------
    # 2) domain_checks index for history queries (watchlist UI)
    # ---------------------------------------------------------
    if await _table_exists(conn, "domain_checks"):
        await conn.execute(
            text(
                "CREATE INDEX IF NOT EXISTS ix_domain_checks_domain_id_checked_at "
                "ON domain_checks(domain_id, checked_at)"
            )
        )

    # ---------------------------------------------------
    # 3) tld_prices composite index for trend computations
    # ---------------------------------------------------
    if await _table_exists(conn, "tld_prices"):
        await conn.execute(
            text(
                "CREATE INDEX IF NOT EXISTS ix_tld_prices_tld_registrar_recorded_at "
                "ON tld_prices(tld, registrar, recorded_at)"
            )
        )

    # ----------------------------------------------------
    # 4) domain_listings pounce_score index (market sorting)
    # ----------------------------------------------------
    if await _table_exists(conn, "domain_listings"):
        await conn.execute(
            text(
                "CREATE INDEX IF NOT EXISTS ix_domain_listings_pounce_score "
                "ON domain_listings(pounce_score)"
            )
        )

    # ----------------------------------------------------
    # 5) Yield tables indexes
    # ----------------------------------------------------
    if await _table_exists(conn, "yield_domains"):
        await conn.execute(
            text(
                "CREATE INDEX IF NOT EXISTS ix_yield_domains_user_status "
                "ON yield_domains(user_id, status)"
            )
        )
        await conn.execute(
            text(
                "CREATE INDEX IF NOT EXISTS ix_yield_domains_domain "
                "ON yield_domains(domain)"
            )
        )

    if await _table_exists(conn, "yield_transactions"):
        await conn.execute(
            text(
                "CREATE INDEX IF NOT EXISTS ix_yield_tx_domain_created "
                "ON yield_transactions(yield_domain_id, created_at)"
            )
        )
        await conn.execute(
            text(
                "CREATE INDEX IF NOT EXISTS ix_yield_tx_status_created "
                "ON yield_transactions(status, created_at)"
            )
        )

    if await _table_exists(conn, "yield_payouts"):
        await conn.execute(
            text(
                "CREATE INDEX IF NOT EXISTS ix_yield_payouts_user_status "
                "ON yield_payouts(user_id, status)"
            )
        )

    # ----------------------------------------------------
    # 6) User referral tracking columns
    # ----------------------------------------------------
    if await _table_exists(conn, "users"):
        if not await _has_column(conn, "users", "referred_by_user_id"):
            logger.info("DB migrations: adding column users.referred_by_user_id")
            await conn.execute(text("ALTER TABLE users ADD COLUMN referred_by_user_id INTEGER"))
        if not await _has_column(conn, "users", "referred_by_domain"):
            logger.info("DB migrations: adding column users.referred_by_domain")
            await conn.execute(text("ALTER TABLE users ADD COLUMN referred_by_domain VARCHAR(255)"))
        if not await _has_column(conn, "users", "referral_code"):
            logger.info("DB migrations: adding column users.referral_code")
            await conn.execute(text("ALTER TABLE users ADD COLUMN referral_code VARCHAR(100)"))

    logger.info("DB migrations: done")