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