spb/pdb-api

1"""Accès DuckDB en lecture seule + garde-fous pour le bac à sable SQL."""2from __future__ import annotations34import os5import re6import threading7import time8from pathlib import Path910import duckdb11import numpy as np1213ROOT = Path(__file__).resolve().parent.parent14DB_PATH = os.environ.get("SPID_DB", str(ROOT / "data" / "spid.duckdb"))15QUERY_TIMEOUT_S = float(os.environ.get("SQL_TIMEOUT_S", "20"))16MAX_ROWS = int(os.environ.get("SQL_MAX_ROWS", "5000"))1718TABLES = ["projects", "project_mentions", "project_timeline", "kg_nodes", "kg_edges", "project_embeddings", "spid_sections"]1920_ALIAS = {"yr": "year", "loc": "location", "lbl": "label"}21_lock = threading.Lock()22_conn: duckdb.DuckDBPyConnection | None = None23_schema_cache: list[dict] | None = None24_emb_cache: dict | None = None252627def connect() -> duckdb.DuckDBPyConnection:28    global _conn29    with _lock:30        if _conn is None:31            if not os.path.exists(DB_PATH):32                raise RuntimeError(f"Base introuvable : {DB_PATH}")33            _conn = duckdb.connect(34                DB_PATH,35                read_only=True,36                config={37                    "enable_external_access": "false",   # pas de read_parquet/read_csv sur le disque38                    "threads": os.environ.get("DUCKDB_THREADS", "4"),39                    "memory_limit": os.environ.get("DUCKDB_MEM", "4GB"),40                },41            )42            try:43                _conn.execute("SET lock_configuration = true")44            except Exception:45                pass46        return _conn474849def cursor() -> duckdb.DuckDBPyConnection:50    """Connexion dupliquée (thread-safe) sur la connexion principale."""51    return connect().cursor()525354def rows(sql: str, params: list | tuple = ()) -> list[dict]:55    cur = cursor()56    try:57        res = cur.execute(sql, list(params))58        cols = [_ALIAS.get(d[0], d[0]) for d in res.description]  # alias courts : year/location/label sont réservés en DuckDB59        return [dict(zip(cols, _jsonable(r))) for r in res.fetchall()]60    finally:61        cur.close()626364def one(sql: str, params: list | tuple = ()) -> dict | None:65    r = rows(sql, params)66    return r[0] if r else None676869def scalar(sql: str, params: list | tuple = ()):70    cur = cursor()71    try:72        return cur.execute(sql, list(params)).fetchone()[0]73    finally:74        cur.close()757677def _jsonable(row):78    out = []79    for v in row:80        if hasattr(v, "isoformat"):81            v = v.isoformat()82        elif isinstance(v, (np.floating,)):83            v = float(v)84        elif isinstance(v, (np.integer,)):85            v = int(v)86        elif isinstance(v, float) and (v != v):  # NaN87            v = None88        out.append(v)89    return out909192# ---------------------------------------------------------------- schéma93def schema() -> list[dict]:94    global _schema_cache95    if _schema_cache is None:96        cur = cursor()97        try:98            out = []99            for t in TABLES:100                cols = cur.execute(101                    "select column_name, data_type from information_schema.columns where table_name=? order by ordinal_position", [t]102                ).fetchall()103                n = cur.execute(f'select count(*) from "{t}"').fetchone()[0]104                out.append({"table": t, "rows": n, "columns": [{"name": c, "type": d} for c, d in cols]})105            _schema_cache = out106        finally:107            cur.close()108    return _schema_cache109110111# ---------------------------------------------------------------- bac à sable SQL112_FORBIDDEN = re.compile(113    r"\b(attach|detach|copy|export|import|install|load|pragma|create|insert|update|delete|drop|alter|call|set|reset|"114    r"checkpoint|vacuum|force|begin|commit|rollback|grant|revoke|use|read_\w+|glob|getenv|current_setting|"115    r"sqlite_\w+|postgres_\w+|http\w*|write_\w+|list_files|duckdb_secrets|secrets?|parquet_\w+|iceberg_\w+|delta_\w+|"116    r"json_extract_path_text|read_json\w*|read_text|read_blob)\b",117    re.IGNORECASE,118)119120121def _strip_comments(sql: str) -> str:122    sql = re.sub(r"/\*.*?\*/", " ", sql, flags=re.S)123    sql = re.sub(r"--[^\n]*", " ", sql)124    return sql.strip()125126127def guard_sql(sql: str) -> str:128    s = _strip_comments(sql).rstrip(";").strip()129    if not s:130        raise ValueError("Requête vide.")131    if ";" in s:132        raise ValueError("Une seule instruction à la fois.")133    if not re.match(r"^(select|with|from|describe|show|summarize|explain)\b", s, re.IGNORECASE):134        raise ValueError("Seules les requêtes de lecture (SELECT / WITH / DESCRIBE / SUMMARIZE / EXPLAIN) sont autorisées.")135    m = _FORBIDDEN.search(s)136    if m:137        raise ValueError(f"Mot-clé interdit dans le bac à sable : {m.group(0)}")138    return s139140141def run_sql(sql: str, limit: int = 500) -> dict:142    """Exécute une requête de lecture avec délai maximal et plafond de lignes."""143    s = guard_sql(sql)144    limit = max(1, min(int(limit), MAX_ROWS))145    wrapped = s if re.match(r"^(describe|show|summarize|explain)\b", s, re.IGNORECASE) else f"select * from ({s}) as _q limit {limit + 1}"146    cur = cursor()147    result: dict = {}148    err: list[Exception] = []149150    def work():151        try:152            t0 = time.perf_counter()153            res = cur.execute(wrapped)154            cols = [d[0] for d in res.description]155            data = res.fetchall()156            result["columns"] = cols157            result["truncated"] = len(data) > limit158            result["rows"] = [_jsonable(r) for r in data[:limit]]159            result["elapsed_ms"] = round((time.perf_counter() - t0) * 1000, 1)160        except Exception as e:  # noqa: BLE001161            err.append(e)162163    th = threading.Thread(target=work, daemon=True)164    th.start()165    th.join(QUERY_TIMEOUT_S)166    if th.is_alive():167        try:168            cur.interrupt()169        except Exception:170            pass171        th.join(2)172        cur.close()173        raise TimeoutError(f"Requête interrompue après {QUERY_TIMEOUT_S:.0f} s.")174    cur.close()175    if err:176        raise ValueError(str(err[0]).split("\n")[0][:400])177    result["sql"] = s178    result["limit"] = limit179    return result180181182# ---------------------------------------------------------------- embeddings (similarité entre projets)183def embeddings() -> dict:184    global _emb_cache185    if _emb_cache is None:186        cur = cursor()187        try:188            data = cur.execute("select project_id, vector from project_embeddings").fetchall()189        finally:190            cur.close()191        ids = [d[0] for d in data]192        mat = np.asarray([d[1] for d in data], dtype=np.float32)193        norms = np.linalg.norm(mat, axis=1, keepdims=True)194        norms[norms == 0] = 1195        mat = mat / norms196        _emb_cache = {"ids": ids, "index": {pid: i for i, pid in enumerate(ids)}, "mat": mat}197    return _emb_cache198199200def similar_ids(project_id: str, k: int = 10) -> list[tuple[str, float]]:201    e = embeddings()202    i = e["index"].get(project_id)203    if i is None:204        return []205    sims = e["mat"] @ e["mat"][i]206    order = np.argsort(-sims)207    out = []208    for j in order:209        if j == i:210            continue211        out.append((e["ids"][j], float(sims[j])))212        if len(out) >= k:213            break214    return out215216217def nearest_to_vector(vec: np.ndarray, k: int = 20) -> list[tuple[str, float]]:218    e = embeddings()219    v = np.asarray(vec, dtype=np.float32)220    v = v / (np.linalg.norm(v) or 1)221    sims = e["mat"] @ v222    order = np.argsort(-sims)[:k]223    return [(e["ids"][j], float(sims[j])) for j in order]224