""" É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