fn_registry/python/functions/datascience/detect_declared_keys_duckdb.py

"""detect_declared_keys_duckdb — lee las claves DECLARADAS de un schema DuckDB.

Funcion impura: lee de disco a traves de la primitiva read-only del grupo
`duckdb` (duckdb_query_readonly). Pertenece al grupo de capacidad `eda`
(relaciones de clave): a diferencia de infer_fk_containment_duckdb, que INFIERE
FOREIGN KEYs candidatas por containment de valores, esta funcion devuelve las
constraints REALES que el schema ha declarado (PRIMARY KEY / FOREIGN KEY /
UNIQUE) leyendo la table function `duckdb_constraints()`.

Es la pieza del capitulo RELACIONES de AutomaticEDA que muestra las relaciones de
clave reales cuando existen — frente a la inferencia, que se usa cuando el schema
no las declaro.

Estilo dict-no-throw del grupo duckdb: nunca lanza; captura cualquier error y
devuelve {status:'error', error:str}.
"""

from infra import duckdb_query_readonly


def _as_list(value) -> list:
    """Normaliza el valor de una columna LIST de DuckDB a una lista de strings.

    En DuckDB 1.5.2, `constraint_column_names` y `referenced_column_names` llegan
    ya como listas Python a traves de duckdb_query_readonly. Este helper es solo
    una red de seguridad: si por cualquier motivo llegara como string (p.ej. la
    representacion `[id, customer_id]`), la parsea de forma defensiva.
    """
    if value is None:
        return []
    if isinstance(value, (list, tuple)):
        return [str(v) for v in value]
    if isinstance(value, str):
        s = value.strip()
        if s.startswith("[") and s.endswith("]"):
            s = s[1:-1]
        if not s.strip():
            return []
        return [
            part.strip().strip("'\"")
            for part in s.split(",")
            if part.strip().strip("'\"")
        ]
    return [str(value)]


def detect_declared_keys_duckdb(db_path: str, table: str = None) -> dict:
    """Detecta las claves PRIMARY KEY / FOREIGN KEY / UNIQUE declaradas en DuckDB.

    Lee la table function `duckdb_constraints()` y extrae solo las constraints de
    clave (PRIMARY KEY, FOREIGN KEY, UNIQUE), ignorando NOT NULL y CHECK.

    Args:
        db_path: ruta al archivo DuckDB. Debe existir (lectura read-only; no se
            crea). Un path inexistente devuelve {status:'error', ...} sin lanzar.
        table: si se pasa, filtra los resultados a esa tabla: incluye PRIMARY KEY
            y UNIQUE cuya tabla sea `table`, y FOREIGN KEY cuya tabla ORIGEN sea
            `table`. None (default) devuelve los constraints de todas las tablas.
            La comparacion de nombres es case-sensitive (tal cual los devuelve
            DuckDB).

    Returns:
        dict dict-no-throw. En exito:
            {status:'ok',
             primary_keys:[{table:str, columns:[str, ...]}, ...],
             foreign_keys:[{table:str, columns:[str, ...],
                            referenced_table:str,
                            referenced_columns:[str, ...]}, ...],
             unique:[{table:str, columns:[str, ...]}, ...],
             tables:[str, ...]}   # tablas (origen) con algun PK/FK/UNIQUE emitido
        En error (sin lanzar): {status:'error', error:str}.
    """
    try:
        sql = (
            "SELECT table_name, constraint_type, constraint_column_names, "
            "referenced_table, referenced_column_names FROM duckdb_constraints()"
        )
        res = duckdb_query_readonly(db_path, sql)
        if res["status"] != "ok":
            return {"status": "error", "error": res["error"]}

        primary_keys = []
        foreign_keys = []
        unique = []
        tables = set()

        for row in res["rows"]:
            ctype = row["constraint_type"]
            tname = row["table_name"]

            # Filtro por tabla origen: para PK/FK/UNIQUE el dueño del constraint es
            # `table_name`. Una FK se atribuye a su tabla origen (no a la
            # referenciada), igual que el filtro pide.
            if table is not None and tname != table:
                continue

            cols = _as_list(row["constraint_column_names"])

            if ctype == "PRIMARY KEY":
                primary_keys.append({"table": tname, "columns": cols})
                tables.add(tname)
            elif ctype == "UNIQUE":
                unique.append({"table": tname, "columns": cols})
                tables.add(tname)
            elif ctype == "FOREIGN KEY":
                foreign_keys.append(
                    {
                        "table": tname,
                        "columns": cols,
                        "referenced_table": row["referenced_table"],
                        "referenced_columns": _as_list(
                            row["referenced_column_names"]
                        ),
                    }
                )
                tables.add(tname)
            # NOT NULL y CHECK se ignoran: no son relaciones de clave.

        return {
            "status": "ok",
            "primary_keys": primary_keys,
            "foreign_keys": foreign_keys,
            "unique": unique,
            "tables": sorted(tables),
        }
    except Exception as e:  # noqa: BLE001
        return {"status": "error", "error": str(e)}