fn_registry/python/functions/datascience/build_join_graph.py

"""Construye un grafo de relaciones inter-tabla a partir de FK candidatas.

Toma la lista `fk_candidates` (salida de infer_fk_containment_duckdb) y produce un
grafo de relaciones: nodos (tablas) con grados y rol inferido (fact/dimension/
bridge/standalone), aristas (una por FK), un diagrama Mermaid pegable y la lista
de hubs (tablas con mayor out_degree, candidatas a tabla de hechos / star schema).

Funcion pura: lista de dicts -> dict de grafo. Sin I/O ni dependencias externas.
"""


def _mermaid_id(name: str) -> str:
    """Sanea un nombre de tabla para usarlo como identificador Mermaid.

    Mermaid no admite espacios, guiones ni puntos en los IDs de nodo. Se sustituyen
    por guion bajo. El nombre original se conserva como etiqueta visible del nodo.
    """
    safe = []
    for ch in str(name):
        safe.append(ch if (ch.isalnum() or ch == "_") else "_")
    out = "".join(safe)
    if not out:
        out = "node"
    if out[0].isdigit():
        out = "t_" + out
    return out


def build_join_graph(fk_candidates: list, tables: list = None) -> dict:
    """Construye un grafo de relaciones inter-tabla desde FK candidatas.

    Args:
        fk_candidates: lista de dicts, cada uno una FK candidata con al menos
            las claves from_table, from_col, to_table, to_col, inclusion,
            cardinality. Claves ausentes se toleran con .get(...). Suele ser la
            salida `fk_candidates` de infer_fk_containment_duckdb.
        tables: lista opcional de nombres de TODAS las tablas. Sirve para incluir
            como nodos aislados (role "standalone") las tablas que no aparecen en
            ninguna FK. Si es None, los nodos se derivan solo de las aristas.

    Returns:
        dict con:
          - nodes: list[dict] con table, out_degree, in_degree, role
                   (role: "fact" | "dimension" | "bridge" | "standalone").
          - edges: list[dict] con from_table, from_col, to_table, to_col,
                   inclusion, cardinality (una por FK valida).
          - mermaid: str con un diagrama `graph LR` pegable en un bloque
                     ```mermaid, una arista por FK etiquetada con las columnas.
          - hubs: list[str] de tablas con mayor out_degree (>0), ordenadas por
                  out_degree descendente. Candidatas a tabla de hechos.
    """
    fk_candidates = fk_candidates or []

    out_degree: dict = {}
    in_degree: dict = {}
    node_order: list = []

    def _ensure(name) -> None:
        if name is None:
            return
        if name not in out_degree:
            out_degree[name] = 0
            in_degree[name] = 0
            node_order.append(name)

    # Sembrar nodos aislados si se pasaron todas las tablas.
    for t in tables or []:
        _ensure(t)

    edges: list = []
    for fk in fk_candidates:
        if not isinstance(fk, dict):
            continue
        ft = fk.get("from_table")
        tt = fk.get("to_table")
        if ft is None or tt is None:
            continue
        _ensure(ft)
        _ensure(tt)
        out_degree[ft] += 1
        in_degree[tt] += 1
        edges.append(
            {
                "from_table": ft,
                "from_col": fk.get("from_col"),
                "to_table": tt,
                "to_col": fk.get("to_col"),
                "inclusion": fk.get("inclusion"),
                "cardinality": fk.get("cardinality"),
            }
        )

    nodes: list = []
    for name in node_order:
        od = out_degree[name]
        ind = in_degree[name]
        if od == 0 and ind == 0:
            role = "standalone"
        elif od > 0 and ind == 0:
            # Apunta a otras tablas pero nadie le apunta: tabla de hechos.
            role = "fact"
        elif od == 0 and ind > 0:
            # Solo recibe referencias: tabla de dimension / maestra.
            role = "dimension"
        else:
            # Apunta y recibe: tabla puente / asociativa.
            role = "bridge"
        nodes.append(
            {
                "table": name,
                "out_degree": od,
                "in_degree": ind,
                "role": role,
            }
        )

    max_out = max((n["out_degree"] for n in nodes), default=0)
    hubs: list = []
    if max_out > 0:
        hubs = [
            n["table"]
            for n in sorted(
                (n for n in nodes if n["out_degree"] > 0),
                key=lambda n: n["out_degree"],
                reverse=True,
            )
        ]

    mermaid = _build_mermaid(nodes, edges)

    return {"nodes": nodes, "edges": edges, "mermaid": mermaid, "hubs": hubs}


def _build_mermaid(nodes: list, edges: list) -> str:
    """Renderiza el grafo como un diagrama Mermaid `graph LR` pegable.

    Una arista por FK, etiquetada con `from_col->to_col`. Los nodos aislados se
    declaran sueltos para que aparezcan en el diagrama. Si no hay nodos ni
    aristas, devuelve un diagrama minimo valido con una nota.
    """
    lines = ["graph LR"]

    if not nodes and not edges:
        lines.append("    empty[No relations]")
        return "\n".join(lines)

    # Declarar nodos aislados (sin ninguna arista) para que se rendericen.
    connected = set()
    for e in edges:
        connected.add(e["from_table"])
        connected.add(e["to_table"])
    for n in nodes:
        name = n["table"]
        if name not in connected:
            nid = _mermaid_id(name)
            lines.append(f'    {nid}["{name}"]')

    for e in edges:
        ft = e["from_table"]
        tt = e["to_table"]
        fc = e.get("from_col")
        tc = e.get("to_col")
        label = f"{fc}->{tc}" if (fc is not None and tc is not None) else ""
        fid = _mermaid_id(ft)
        tid = _mermaid_id(tt)
        if label:
            lines.append(f'    {fid}["{ft}"] -->|{label}| {tid}["{tt}"]')
        else:
            lines.append(f'    {fid}["{ft}"] --> {tid}["{tt}"]')

    return "\n".join(lines)