metalfrom.eu/apps/crawler/src/db.py

"""
Écriture directe en DB. Toutes les opérations passent par des upserts idempotents.
"""
import json
import logging
from contextlib import contextmanager
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional

import psycopg2
import psycopg2.extras

from .config import DATABASE_URL

log = logging.getLogger(__name__)


@contextmanager
def get_conn():
    conn = psycopg2.connect(DATABASE_URL)
    try:
        yield conn
        conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        conn.close()


def now_utc() -> datetime:
    return datetime.now(timezone.utc)


# ------------------------------------------------------------------
# Bands
# ------------------------------------------------------------------

def upsert_bands(bands: List[Dict[str, Any]]) -> Dict[str, int]:
    """
    Upsert en bulk. Préserve les champs existants si la nouvelle valeur est None.
    Retourne {"inserted": N, "updated": N}.
    """
    if not bands:
        return {"inserted": 0, "updated": 0}

    ts = now_utc()
    rows = []
    for b in bands:
        rows.append((
            b["ma_id"],
            b.get("name") or "",
            b.get("country"),
            b.get("location_text"),
            b.get("status"),
            b.get("genre"),
            b.get("formed_year"),
            b.get("themes"),
            b.get("enriched", False),
            b.get("crawled_at"),
            b.get("crawled_hash"),
            b.get("ma_created_at"),
            b.get("ma_modified_at"),
            ts,   # first_seen_at (ignoré si déjà set)
            ts,   # created_at (ignoré si déjà set)
            psycopg2.extras.Json(b.get("data") or {}),
        ))

    sql = """
        INSERT INTO bands
          (ma_id, name, country, location_text, status, genre, formed_year,
           themes, enriched, crawled_at, crawled_hash, ma_created_at, ma_modified_at,
           first_seen_at, created_at, data)
        VALUES %s
        ON CONFLICT (ma_id) DO UPDATE SET
          name          = COALESCE(EXCLUDED.name, bands.name),
          country       = COALESCE(EXCLUDED.country, bands.country),
          location_text = COALESCE(EXCLUDED.location_text, bands.location_text),
          status        = COALESCE(EXCLUDED.status, bands.status),
          genre         = COALESCE(EXCLUDED.genre, bands.genre),
          formed_year   = COALESCE(EXCLUDED.formed_year, bands.formed_year),
          themes        = COALESCE(EXCLUDED.themes, bands.themes),
          enriched      = GREATEST(EXCLUDED.enriched, bands.enriched),
          crawled_at    = COALESCE(EXCLUDED.crawled_at, bands.crawled_at),
          crawled_hash  = COALESCE(EXCLUDED.crawled_hash, bands.crawled_hash),
          ma_created_at = COALESCE(EXCLUDED.ma_created_at, bands.ma_created_at),
          ma_modified_at= COALESCE(EXCLUDED.ma_modified_at, bands.ma_modified_at),
          data          = bands.data || EXCLUDED.data
        RETURNING (xmax = 0) AS was_inserted
    """

    with get_conn() as conn:
        with conn.cursor() as cur:
            results = psycopg2.extras.execute_values(cur, sql, rows, fetch=True)
            inserted = sum(1 for r in results if r[0])
            updated  = len(results) - inserted

    log.info(f"[db] upsert_bands: {inserted} inserted, {updated} updated")
    return {"inserted": inserted, "updated": updated}


def upsert_band_enriched(ma_id: int, data: Dict[str, Any], html_hash: str) -> bool:
    """Met à jour les champs d'enrichissement d'un band."""
    sql = """
        UPDATE bands SET
          status        = COALESCE(%s, status),
          genre         = COALESCE(%s, genre),
          themes        = COALESCE(%s, themes),
          formed_year   = COALESCE(%s, formed_year),
          data          = data || %s::jsonb,
          enriched      = true,
          crawled_at    = %s,
          crawled_hash  = %s,
          ma_created_at = COALESCE(%s, ma_created_at),
          ma_modified_at= COALESCE(%s, ma_modified_at)
        WHERE ma_id = %s
    """
    ts = now_utc()
    with get_conn() as conn:
        with conn.cursor() as cur:
            cur.execute(sql, (
                data.get("status"),
                data.get("genre"),
                data.get("themes"),
                _parse_year(data.get("formed_in")),
                json.dumps({"band_page": data, "enriched_at": ts.isoformat()}),
                ts,
                html_hash,
                data.get("ma_created_at"),
                data.get("ma_modified_at"),
                ma_id,
            ))
            return cur.rowcount > 0


def get_bands_to_enrich(country: Optional[str] = None, limit: int = 50) -> List[Dict]:
    """
    File de priorité pour l'enrichissement :
      1. Nouveaux bands (band_page absent) — jamais enrichis
      2. Modifiés depuis le dernier enrichissement (updated_at > crawled_at + 1 min)
      3. Héritage ancien scraper (crawled_at IS NULL, band_page présent)
      4. Stale (enrichis il y a > 30 jours par le système actuel)
    """
    sql = """
        SELECT ma_id, data->>'url' AS url, name, country
        FROM bands
        WHERE data->>'url' IS NOT NULL
          AND (
            data->'band_page' IS NULL
            OR crawled_at IS NULL
            OR (crawled_at IS NOT NULL AND updated_at > crawled_at + interval '1 minute')
            OR crawled_at < now() - interval '30 days'
          )
    """
    params: list = []
    if country:
        sql += " AND country = %s"
        params.append(country)
    sql += """
        ORDER BY
          CASE
            WHEN data->'band_page' IS NULL THEN 1
            WHEN crawled_at IS NOT NULL AND updated_at > crawled_at + interval '1 minute' THEN 2
            WHEN crawled_at IS NULL THEN 3
            ELSE 4
          END ASC,
          updated_at DESC NULLS LAST
        LIMIT %s
    """
    params.append(limit)

    with get_conn() as conn:
        with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
            cur.execute(sql, params)
            return [dict(r) for r in cur.fetchall()]


# ------------------------------------------------------------------
# Logs (visibles dans le dashboard admin)
# ------------------------------------------------------------------

def log_event(level: str, message: str, run_id: Optional[int] = None, ma_id: Optional[int] = None):
    """Écrit une ligne de log en DB pour le dashboard admin (n'interrompt jamais le crawl)."""
    try:
        with get_conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    "INSERT INTO crawl_log (run_id, level, message, ma_id) VALUES (%s, %s, %s, %s)",
                    (run_id, level, message[:2000], ma_id),
                )
    except Exception as e:
        log.warning(f"[db] log_event failed: {e}")


# ------------------------------------------------------------------
# Crawl run tracking
# ------------------------------------------------------------------

def update_crawl_run_progress(run_id: int, stats: Dict[str, int]):
    """Mise à jour des compteurs d'un run en cours (progression live)."""
    try:
        with get_conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """UPDATE crawl_run SET
                       bands_seen=%s, bands_new=%s, bands_updated=%s, bands_enriched=%s
                       WHERE id=%s AND status='running'""",
                    (stats.get("seen", 0), stats.get("new", 0),
                     stats.get("updated", 0), stats.get("enriched", 0), run_id),
                )
    except Exception as e:
        log.warning(f"[db] update_crawl_run_progress failed: {e}")


def start_crawl_run(run_type: str, countries: Optional[List[str]] = None) -> int:
    sql = "INSERT INTO crawl_run (run_type, countries) VALUES (%s, %s) RETURNING id"
    with get_conn() as conn:
        with conn.cursor() as cur:
            cur.execute(sql, (run_type, countries))
            return cur.fetchone()[0]


def finish_crawl_run(run_id: int, stats: Dict[str, int], error: Optional[str] = None):
    status = "error" if error else "done"
    sql = """
        UPDATE crawl_run SET
          status = %s, finished_at = now(),
          bands_seen = %s, bands_new = %s, bands_updated = %s, bands_enriched = %s,
          error = %s
        WHERE id = %s
    """
    with get_conn() as conn:
        with conn.cursor() as cur:
            cur.execute(sql, (
                status,
                stats.get("seen", 0), stats.get("new", 0),
                stats.get("updated", 0), stats.get("enriched", 0),
                error, run_id,
            ))


# ------------------------------------------------------------------
# Checkpoints
# ------------------------------------------------------------------

def get_checkpoint(key: str) -> Optional[str]:
    with get_conn() as conn:
        with conn.cursor() as cur:
            cur.execute("SELECT value FROM crawl_checkpoint WHERE key = %s", (key,))
            row = cur.fetchone()
            return row[0] if row else None


def set_checkpoint(key: str, value: str):
    sql = """
        INSERT INTO crawl_checkpoint (key, value, updated_at)
        VALUES (%s, %s, now())
        ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value, updated_at = now()
    """
    with get_conn() as conn:
        with conn.cursor() as cur:
            cur.execute(sql, (key, value))


# ------------------------------------------------------------------
# Helpers
# ------------------------------------------------------------------

def _parse_year(s: Optional[str]) -> Optional[int]:
    if not s:
        return None
    m = __import__("re").search(r"\b(1[89]\d\d|20\d\d)\b", str(s))
    return int(m.group(1)) if m else None