hechakuaa_bot/database.py

"""
database.py — Persistencia de tenants y preferencias.

SQLite vía aiosqlite. En v1 esto se reemplaza por Postgres con el mismo interface público.

Schema:
- tenants: un registro por abogado registrado.
  Las credenciales van cifradas con vault.py (Fernet derivada por chat_id).
- Migraciones: directorio migrations/*.sql, ordenadas por nombre.
  schema_migrations lleva registro de las ya aplicadas.
"""
from __future__ import annotations

from datetime import datetime, timezone
from pathlib import Path

import aiosqlite
import structlog

logger = structlog.get_logger(__name__)

_conn: aiosqlite.Connection | None = None


def get_conn() -> aiosqlite.Connection:
    """Retorna la conexión activa. Llamar solo después de run_migrations()."""
    return _assert_conn()


def _now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


async def run_migrations(db_path: Path) -> None:
    """
    Aplica migraciones pendientes desde migrations/*.sql en orden alfabético.
    Idempotente: versiones ya aplicadas se saltean.
    Reemplaza init_db() — llamar una sola vez al arrancar.
    """
    global _conn
    db_path.parent.mkdir(parents=True, exist_ok=True)
    _conn = await aiosqlite.connect(str(db_path))
    _conn.row_factory = aiosqlite.Row

    await _conn.execute("""
        CREATE TABLE IF NOT EXISTS schema_migrations (
            version    TEXT PRIMARY KEY,
            applied_at TEXT NOT NULL
        )
    """)
    await _conn.commit()

    async with _conn.execute("SELECT version FROM schema_migrations") as cur:
        applied = {row["version"] for row in await cur.fetchall()}

    migrations_dir = Path(__file__).parent / "migrations"
    for migration_file in sorted(migrations_dir.glob("*.sql")):
        version = migration_file.stem
        if version in applied:
            logger.debug("migration_ya_aplicada", version=version)
            continue

        sql = migration_file.read_text(encoding="utf-8")
        for stmt in sql.split(";"):
            stmt = stmt.strip()
            # Saltear líneas de comentario puras y sentencias vacías
            if stmt and not all(line.startswith("--") for line in stmt.splitlines() if line.strip()):
                await _conn.execute(stmt)

        await _conn.execute(
            "INSERT INTO schema_migrations (version, applied_at) VALUES (?, ?)",
            (version, _now_iso()),
        )
        await _conn.commit()
        logger.info("migration_aplicada", version=version)

    logger.info("database_iniciada", path=str(db_path))


async def init_db(db_path: Path) -> None:
    """Alias de run_migrations() para compatibilidad con scripts existentes."""
    await run_migrations(db_path)


def _assert_conn() -> aiosqlite.Connection:
    if _conn is None:
        raise RuntimeError("DB no inicializada — llamar a init_db() antes de usar database.")
    return _conn


async def get_tenant(chat_id: int) -> dict | None:
    """Retorna el tenant o None si no existe."""
    conn = _assert_conn()
    async with conn.execute(
        "SELECT * FROM tenants WHERE chat_id = ?", (chat_id,)
    ) as cur:
        row = await cur.fetchone()
        return dict(row) if row else None


async def save_tenant(chat_id: int, **fields: object) -> None:
    """
    Crear o actualizar un tenant. Los kwargs son los campos de la tabla.
    Hace upsert: inserta si no existe, actualiza si ya existe.
    """
    conn = _assert_conn()
    fields["last_active"] = _now_iso()
    existing = await get_tenant(chat_id)

    if existing is None:
        fields.setdefault("created_at", _now_iso())
        columns = ["chat_id", *fields.keys()]
        placeholders = ", ".join("?" * len(columns))
        values = [chat_id, *fields.values()]
        await conn.execute(
            f"INSERT INTO tenants ({', '.join(columns)}) VALUES ({placeholders})",
            values,
        )
        logger.info("tenant_creado", chat_id=chat_id)
    else:
        set_clause = ", ".join(f"{k} = ?" for k in fields)
        values_upd = [*fields.values(), chat_id]
        await conn.execute(
            f"UPDATE tenants SET {set_clause} WHERE chat_id = ?",
            values_upd,
        )
        logger.info("tenant_actualizado", chat_id=chat_id, campos=list(fields.keys()))

    await conn.commit()


async def list_active_tenants() -> list[dict]:
    """Retorna todos los tenants con estado='activo'."""
    conn = _assert_conn()
    async with conn.execute(
        "SELECT * FROM tenants WHERE estado = 'activo'"
    ) as cur:
        rows = await cur.fetchall()
        return [dict(r) for r in rows]


async def delete_tenant(chat_id: int) -> None:
    """Elimina el tenant y sus datos de la DB."""
    conn = _assert_conn()
    await conn.execute("DELETE FROM tenants WHERE chat_id = ?", (chat_id,))
    await conn.commit()
    logger.info("tenant_eliminado", chat_id=chat_id)


# ── Invite codes ──────────────────────────────────────────────────────────────

async def crear_invite_code(codigo: str) -> None:
    """Inserta un código de invitación de un solo uso."""
    conn = _assert_conn()
    await conn.execute(
        "INSERT INTO invite_codes (codigo, creado_at) VALUES (?, ?)",
        (codigo, _now_iso()),
    )
    await conn.commit()
    logger.info("invite_code_creado", codigo=codigo[:4] + "****")


async def verificar_y_consumir_invite_code(codigo: str, chat_id: int) -> bool:
    """
    Verifica y consume un código de invitación en una sola operación atómica.

    Usa UPDATE con condición AND usado=0 en vez de SELECT+UPDATE para eliminar la
    race condition TOCTOU: dos requests concurrentes con el mismo código solo
    pueden actualizar la fila una vez — el segundo UPDATE afecta 0 rows y retorna False.
    """
    conn = _assert_conn()
    cur = await conn.execute(
        "UPDATE invite_codes SET usado = 1, usado_por = ?, usado_at = ? "
        "WHERE codigo = ? AND usado = 0",
        (chat_id, _now_iso(), codigo),
    )
    await conn.commit()
    if cur.rowcount > 0:
        logger.info("invite_code_consumido", chat_id=chat_id)
        return True
    return False