fn_registry/python/functions/infra/duckdb_upsert.py

"""UPSERT idempotente de filas en una tabla DuckDB con ownership selectivo de columnas.

Funcion impura: abre un archivo DuckDB con `duckdb.connect(db_path)` en modo
lectura-escritura, ejecuta un `INSERT ... ON CONFLICT (key_cols) DO UPDATE SET ...`
fila por fila, hace commit y cierra la conexion en un bloque try/finally. Devuelve
un dict sin lanzar excepciones, siguiendo el estilo del grupo duckdb del registry:
{status:'ok', ...} en exito y {status:'error', error:str} en fallo.

El valor de esta funcion es el "ownership selectivo": al actualizar solo las
columnas indicadas en `update_cols` en caso de conflicto, un re-upsert de la misma
clave NO pisa las columnas que se dejaron fuera de `update_cols`. Asi la DB puede
ser la fuente de verdad de ciertos campos (enriquecidos, anotados a mano, derivados)
mientras un proceso de re-ingest refresca solo los campos que aporta.

Identificadores (tabla y columnas) se validan contra `[A-Za-z_][A-Za-z0-9_]*` antes
de interpolarlos en el SQL (DuckDB no permite parametrizar identificadores); los
valores de las filas siempre van por placeholders `?`.
"""

import re

_IDENT_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")


def _validate_ident(name: str) -> str:
    """Valida que `name` sea un identificador SQL seguro y lo devuelve.

    Acepta solo nombres que casen `[A-Za-z_][A-Za-z0-9_]*`. Lanza ValueError
    para cualquier otro (espacios, comillas, puntos, vacio), que el caller
    convierte en {status:'error'}.
    """
    if not isinstance(name, str) or not _IDENT_RE.match(name):
        raise ValueError(f"identificador invalido: {name!r}")
    return name


def duckdb_upsert(
    db_path: str,
    table: str,
    rows: list,
    key_cols: list,
    update_cols: list = None,
) -> dict:
    """Hace UPSERT idempotente de `rows` en `table`, con ownership selectivo.

    Construye `INSERT INTO <table> (cols) VALUES (?,...) ON CONFLICT (key_cols)
    DO UPDATE SET col=excluded.col, ...` (o `DO NOTHING`) y lo ejecuta fila por
    fila para poder contar inserts vs updates.

    Args:
        db_path: ruta al archivo DuckDB. Se abre en lectura-escritura
            (`duckdb.connect`), por lo que se crea si no existe — pero la tabla
            destino debe existir y tener PRIMARY KEY o UNIQUE en `key_cols`.
        table: nombre de la tabla destino. Validado como identificador SQL.
        rows: lista de dicts, un dict por fila (clave=nombre de columna). El
            esquema de insercion lo fija el conjunto de claves de la PRIMERA fila;
            todas las filas deben tener exactamente las mismas claves o se devuelve
            {status:'error'}. Lista vacia -> {status:'ok', inserted:0, updated:0}.
        key_cols: columnas de la clave de conflicto. Deben tener PRIMARY KEY o
            UNIQUE en la tabla para que ON CONFLICT funcione. Deben estar presentes
            en las claves de las filas.
        update_cols: columnas a actualizar en caso de conflicto.
            None (default) -> todas las columnas de la fila MENOS las key_cols.
            Lista vacia [] -> DO NOTHING (inserta nuevas, no toca existentes).
            Lista con columnas -> DO UPDATE SET solo esas (las no listadas
            conservan su valor previo: ownership selectivo).

    Returns:
        dict. En exito: {status:'ok', inserted:int, updated:int} donde inserted
        cuenta las claves que no existian y updated las que ya existian (para
        update_cols=[] -> DO NOTHING, updated es 0). En error (sin lanzar):
        {status:'error', error:str}.
    """
    conn = None
    try:
        if not isinstance(rows, list):
            raise ValueError("rows debe ser una lista de dicts")
        if not rows:
            return {"status": "ok", "inserted": 0, "updated": 0}

        # Esquema de insercion = claves de la primera fila, en orden estable.
        first_keys = list(rows[0].keys())
        insert_cols = [_validate_ident(c) for c in first_keys]
        insert_set = set(first_keys)

        # Todas las filas deben tener exactamente las mismas claves.
        for i, row in enumerate(rows):
            if not isinstance(row, dict):
                raise ValueError(f"rows[{i}] no es un dict")
            if set(row.keys()) != insert_set:
                raise ValueError(
                    f"rows[{i}] tiene columnas distintas a la primera fila: "
                    f"{sorted(row.keys())} vs {sorted(first_keys)}"
                )

        keys = [_validate_ident(c) for c in key_cols]
        if not keys:
            raise ValueError("key_cols no puede estar vacio")
        for k in keys:
            if k not in insert_set:
                raise ValueError(f"key_col {k!r} no esta en las columnas de las filas")

        # Resolver update_cols.
        if update_cols is None:
            updates = [c for c in insert_cols if c not in keys]
        else:
            updates = [_validate_ident(c) for c in update_cols]
            for u in updates:
                if u not in insert_set:
                    raise ValueError(
                        f"update_col {u!r} no esta en las columnas de las filas"
                    )

        cols_sql = ", ".join(insert_cols)
        placeholders = ", ".join(["?"] * len(insert_cols))
        conflict_sql = ", ".join(keys)

        if updates:
            set_sql = ", ".join(f"{c} = excluded.{c}" for c in updates)
            on_conflict = f"ON CONFLICT ({conflict_sql}) DO UPDATE SET {set_sql}"
        else:
            on_conflict = f"ON CONFLICT ({conflict_sql}) DO NOTHING"

        sql = (
            f"INSERT INTO {table} ({cols_sql}) VALUES ({placeholders}) {on_conflict}"
        )

        conn = __import__("duckdb").connect(db_path)

        # Contamos insert vs update consultando la existencia de la clave antes
        # de ejecutar cada fila. Misma conexion/transaccion, single-writer.
        where_keys = " AND ".join(f"{k} = ?" for k in keys)
        exists_sql = f"SELECT 1 FROM {table} WHERE {where_keys} LIMIT 1"

        inserted = 0
        updated = 0
        for row in rows:
            key_vals = [row[k] for k in keys]
            existed = conn.execute(exists_sql, key_vals).fetchone() is not None

            values = [row[c] for c in insert_cols]
            conn.execute(sql, values)

            if existed:
                updated += 1
            else:
                inserted += 1

        conn.commit()
        return {"status": "ok", "inserted": inserted, "updated": updated}
    except Exception as e:  # noqa: BLE001
        return {"status": "error", "error": str(e)}
    finally:
        if conn is not None:
            conn.close()