fn_registry/python/functions/infra/pg_query.py

"""Ejecuta una query SELECT contra PostgreSQL y devuelve filas como list[dict].

Funcion impura: abre una conexion psycopg2 con el DSN dado, ejecuta el SQL con un
RealDictCursor (cada fila es un dict columna->valor) y devuelve un dict sin lanzar
excepciones, siguiendo el estilo de duckdb_query_readonly del registry:
{status:'ok', ...} en exito y {status:'error', error:str} en fallo. La conexion se
cierra siempre en un bloque try/finally.

Por convencion es de solo lectura: la transaccion se marca read-only
(SET TRANSACTION READ ONLY) para que cualquier escritura accidental falle a nivel
de servidor, y nunca se hace commit (rollback al final). El resultado se trunca a
max_rows para proteger la memoria y marca truncated=True si la query producia mas
filas. Los valores que no son JSON-serializables se convierten a una forma
serializable: date/datetime/time a isoformat(), Decimal a float, bytes/memoryview a
base64 y UUID a str.
"""

import base64
import datetime
import decimal
import uuid


def _to_serializable(value):
    """Convierte un valor de PostgreSQL a una forma JSON-serializable.

    date/datetime/time -> isoformat(), Decimal -> float, bytes/memoryview -> base64
    str, UUID -> str. El resto de valores (int, float, str, bool, None) se devuelven
    sin cambios.
    """
    if value is None:
        return None
    if isinstance(value, (datetime.datetime, datetime.date, datetime.time)):
        return value.isoformat()
    if isinstance(value, decimal.Decimal):
        return float(value)
    if isinstance(value, (bytes, bytearray, memoryview)):
        return base64.b64encode(bytes(value)).decode("ascii")
    if isinstance(value, uuid.UUID):
        return str(value)
    return value


def pg_query(
    dsn: str,
    sql: str,
    params: list = None,
    max_rows: int = 10000,
) -> dict:
    """Ejecuta un SELECT contra PostgreSQL en una transaccion read-only.

    Args:
        dsn: cadena de conexion PostgreSQL, p.ej.
            "postgresql://user:pass@localhost:5433/trends". Un DSN invalido o un
            servidor inalcanzable devuelve {status:'error', ...} (no lanza).
        sql: sentencia SQL a ejecutar. Pensada para SELECT; usa el marcador `%s`
            para parametros posicionales (estilo psycopg2).
        params: lista de parametros posicionales para el SQL, en orden. None
            (default) significa sin parametros. Pasar los valores aqui en vez de
            interpolarlos en el SQL evita inyeccion.
        max_rows: numero maximo de filas a materializar (default 10000). Si la
            query produce mas, se trunca y truncated queda en True.

    Returns:
        dict. En exito: {status:'ok', columns:[str,...], rows:[{col:val, ...}, ...],
        row_count:int, truncated:bool} donde columns es la lista de nombres de
        columna y rows es la lista de filas (cada fila un dict, via RealDictCursor).
        En error (sin lanzar): {status:'error', error:str}.
    """
    try:
        import psycopg2
        from psycopg2 import extras as pg_extras
    except ImportError as exc:  # pragma: no cover - exercised only without dep
        return {
            "status": "error",
            "error": (
                "psycopg2 is required for pg_query; install psycopg2-binary "
                f"({exc})"
            ),
        }

    conn = None
    try:
        conn = psycopg2.connect(dsn)
        # Solo lectura por convencion: cualquier escritura fallara en el servidor.
        conn.set_session(readonly=True, autocommit=False)
        with conn.cursor(cursor_factory=pg_extras.RealDictCursor) as cur:
            cur.execute(sql, params if params is not None else None)

            description = cur.description or []
            columns = [col.name for col in description]

            # Pedimos una fila de mas que max_rows para detectar truncado.
            fetched = cur.fetchmany(max_rows + 1)
            truncated = len(fetched) > max_rows
            if truncated:
                fetched = fetched[:max_rows]

            rows = [
                {key: _to_serializable(val) for key, val in record.items()}
                for record in fetched
            ]

        # Nunca escribimos: cerramos la transaccion con rollback.
        conn.rollback()

        return {
            "status": "ok",
            "columns": columns,
            "rows": rows,
            "row_count": len(rows),
            "truncated": truncated,
        }
    except Exception as e:  # noqa: BLE001
        return {"status": "error", "error": str(e)}
    finally:
        if conn is not None:
            conn.close()