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

"""
Écriture directe en DB. Toutes les opérations passent par des upserts idempotents.
"""
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)
        ))

    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)
        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)
        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($2, status),
          genre         = COALESCE($3, genre),
          themes        = COALESCE($4, themes),
          formed_year   = COALESCE($5, formed_year),
          data          = data || $6::jsonb,
          enriched      = true,
          crawled_at    = $7,
          crawled_hash  = $8,
          ma_created_at = COALESCE($9, ma_created_at),
          ma_modified_at= COALESCE($10, ma_modified_at)
        WHERE ma_id = $1
    """
    ts = now_utc()
    with get_conn() as conn:
        with conn.cursor() as cur:
            import json
            cur.execute(sql, (
                ma_id,
                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"),
            ))
            return cur.rowcount > 0


def get_bands_to_enrich(country: Optional[str] = None, limit: int = 50) -> List[Dict]:
    """Bands qui n'ont jamais été enrichies (data->'band_page' IS NULL)."""
    sql = """
        SELECT ma_id, data->>'url' AS url, name, country
        FROM bands
        WHERE data->'band_page' IS NULL
          AND data->>'url' IS NOT NULL
    """
    params = []
    if country:
        sql += " AND country = %s"
        params.append(country)
    sql += " ORDER BY first_seen_at ASC 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()]


def get_bands_stale(hours: int = 24 * 30, limit: int = 100) -> List[Dict]:
    """Bands enrichies il y a plus de N heures (pour re-crawl périodique)."""
    sql = """
        SELECT ma_id, data->>'url' AS url, name, country
        FROM bands
        WHERE enriched = true
          AND (crawled_at IS NULL OR crawled_at < now() - make_interval(hours => %s))
        ORDER BY crawled_at ASC NULLS FIRST
        LIMIT %s
    """
    with get_conn() as conn:
        with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
            cur.execute(sql, (hours, limit))
            return [dict(r) for r in cur.fetchall()]


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

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