fn_registry/python/functions/pipelines/profile_bq_dataset.py

"""profile_bq_dataset — EDA one-shot de un dataset BigQuery ENTERO (cross-tabla).

Pipeline impuro: perfila un dataset de Google BigQuery COMPLETO con descubrimiento
cross-tabla (relaciones FK inter-tabla + join graph + diccionario de columnas). Es
el analogo a nivel de dataset de `profile_bq_table` (que perfila UNA tabla) y el
adaptador BigQuery de `profile_database` (que perfila una base DuckDB entera),
resuelto por composicion estricta de funciones del registry — sin reimplementar
ni el descubrimiento, ni el perfilado, ni la inferencia de relaciones.

Clave del diseno: materializa CADA tabla del dataset en UN MISMO archivo DuckDB
compartido, para que la inferencia de FK por containment (que necesita todas las
tablas en la misma base) opere cross-tabla. El DuckDB compartido queda como el
artefacto explorable post-EDA (keep_duckdb=True por defecto).

Funciones del registry compuestas (NO se reimplementa su logica):
  - list_bq_dataset_tables : catalogo de tablas/vistas del dataset (descubrimiento).
  - load_bq_table_to_duckdb: materializa cada tabla BigQuery al DuckDB compartido
                             (completa por defecto, muestra si sample_frac; PII
                             seudonimizada por tabla antes de escribir a disco).
  - profile_database       : perfila el DuckDB compartido entero (perfiles por
                             tabla + FK por containment + join graph Mermaid +
                             report markdown/JSON + PDF movil DB-level).
  - build_column_dictionary: diccionario de columnas buscable (tipo semantico, PII)
                             a partir del DatabaseProfile ensamblado.

Modo por defecto = FULL: `sample_frac=None` trae cada tabla entera (preferencia
estandar del usuario: los EDA se corren sobre el total salvo que se pida lo
contrario). El muestreo es opt-in explicito por dataset: `sample_frac=0.05` trae
~5 % de cada tabla; `max_rows` es un tope duro por tabla (0 = sin tope).

Seudonimizacion LOPDGDD/RGPD [POL-MMNSEG-001-1.0]: `pseudonymize_cols` es un dict
`{"tabla": ["col1", "col2"]}` que se aplica por tabla ANTES de materializar.

Estilo dict-no-throw del grupo `eda`: nunca lanza; captura cualquier error y
devuelve {status:'error', error, stage}. Los fallos por tabla individual se
toleran: se anotan en errors[]/tables_skipped[] y se sigue con las demas.
"""

import json
import os
from datetime import datetime, timezone

from datascience import (
    build_column_dictionary,
    list_bq_dataset_tables,
    load_bq_table_to_duckdb,
)
from pipelines.profile_database import profile_database


def profile_bq_dataset(
    project_id: str,
    dataset: str,
    tables: list = None,
    include_views: bool = False,
    sample_frac: float = None,
    max_rows: int = 0,
    pseudonymize_cols: dict = None,
    report_dir: str = "reports",
    duckdb_path: str = "",
    keep_duckdb: bool = True,
    min_inclusion: float = 0.9,
    emit_pdf: bool = True,
    run_llm: bool = False,
) -> dict:
    """EDA one-shot de un dataset BigQuery entero, con descubrimiento cross-tabla.

    Materializa cada tabla del dataset a un DuckDB compartido, lo perfila entero
    con `profile_database` (perfiles por tabla + FK inter-tabla + join graph +
    reports + PDF) y construye el diccionario de columnas del dataset. Por defecto
    perfila TODAS las filas de cada tabla (`sample_frac=None`, modo FULL) y solo
    las BASE TABLE (las vistas se excluyen salvo `include_views=True`).

    Args:
        project_id: proyecto GCP (facturacion + primer segmento del FQN).
        dataset: dataset BigQuery a perfilar.
        tables: lista de NOMBRES de tabla del dataset. None (default) = todas las
            del dataset (filtradas por include_views).
        include_views: si True incluye las VIEW ademas de las BASE TABLE cuando
            tables=None. Default False (solo BASE TABLE, coherente con
            profile_database que salta las VIEWs).
        sample_frac: None (default) = FULL, perfila todas las filas de cada tabla.
            Un float en (0,1) activa el muestreo opt-in por tabla.
        max_rows: tope duro de filas por tabla (LIMIT). 0 (default) = sin tope.
        pseudonymize_cols: dict {"tabla": ["col1", "col2"]} de columnas PII a
            seudonimizar (hash) por tabla ANTES de materializar (LOPDGDD/RGPD).
        report_dir: directorio de salida de los reports + del DuckDB por defecto.
        duckdb_path: ruta del DuckDB compartido. Vacio = report_dir/eda_bq_<dataset>.duckdb.
        keep_duckdb: si True (default) conserva el DuckDB materializado (es el
            artefacto explorable post-EDA). Con False se borra al terminar.
        min_inclusion: umbral de inclusion (0-1) para emitir una FK candidata
            (se pasa a profile_database -> infer_fk_containment_duckdb).
        emit_pdf: si True (default) emite el PDF movil DB-level (se pasa a
            profile_database).
        run_llm: si True (default False) activa la capa LLM interpretativa por
            tabla (se pasa a profile_database -> profile_table; una llamada LLM
            por tabla sobre el perfil agregado, nunca filas crudas).

    Returns:
        dict dict-no-throw con el resultado del pipeline (ver output del .md).
    """
    try:
        # 1) Resolver la lista de tablas a materializar como (nombre, fqn).
        if tables is None:
            lst = list_bq_dataset_tables(
                project_id, dataset, include_views=include_views
            )
            if lst.get("status") != "ok":
                return {
                    "status": "error",
                    "error": lst.get("error", "list_bq_dataset_tables fallo"),
                    "stage": "list_tables",
                }
            table_specs = [
                (t["table"], t.get("fqn") or f"{project_id}.{dataset}.{t['table']}")
                for t in lst.get("tables", [])
            ]
        elif isinstance(tables, list):
            table_specs = [
                (name, f"{project_id}.{dataset}.{name}") for name in tables
            ]
        else:
            return {
                "status": "error",
                "error": "tables debe ser una lista de nombres o None",
                "stage": "list_tables",
            }

        if not table_specs:
            return {
                "status": "error",
                "error": (
                    "el dataset no tiene tablas que perfilar "
                    "(revisa include_views / el nombre del dataset)"
                ),
                "stage": "list_tables",
            }

        # 2) Resolver el DuckDB compartido. Vacio -> report_dir/eda_bq_<dataset>.duckdb.
        os.makedirs(report_dir, exist_ok=True)
        created_default = False
        if not duckdb_path:
            duckdb_path = os.path.join(report_dir, f"eda_bq_{dataset}.duckdb")
            created_default = True
        # Si es el path por defecto y ya existe de una corrida previa, empezar
        # limpio para que no se mezclen tablas de datasets distintos en la base.
        if created_default and os.path.exists(duckdb_path):
            try:
                os.remove(duckdb_path)
            except OSError:
                pass

        # 3) Materializar CADA tabla en el MISMO DuckDB (tolerando fallos).
        pseudo_map = pseudonymize_cols or {}
        loaded_tables = []       # nombres de tabla dentro del DuckDB
        tables_skipped = []
        errors = []
        for name, fqn in table_specs:
            load = load_bq_table_to_duckdb(
                fqn,
                duckdb_path,
                sample_frac=sample_frac,
                max_rows=max_rows,
                project_id=project_id,
                pseudonymize_cols=pseudo_map.get(name),
            )
            if load.get("status") == "ok":
                loaded_tables.append(load["table"])
            else:
                errors.append(
                    {
                        "table": name,
                        "stage": "load",
                        "error": load.get("error", "load fallo"),
                    }
                )
                tables_skipped.append(name)

        if not loaded_tables:
            return {
                "status": "error",
                "error": "ninguna tabla del dataset se pudo materializar",
                "stage": "load",
                "errors": errors,
            }

        # 4) Perfilar el DuckDB compartido entero: perfiles por tabla + FK
        # cross-tabla + join graph + report markdown/JSON + PDF (si emit_pdf).
        prof = profile_database(
            duckdb_path,
            tables=loaded_tables,
            report_dir=report_dir,
            write_report=True,
            min_inclusion=min_inclusion,
            emit_pdf=emit_pdf,
            run_llm=run_llm,
        )
        if prof.get("status") != "ok":
            return {
                "status": "error",
                "error": prof.get("error", "profile_database fallo"),
                "stage": "profile",
            }
        db_profile = prof.get("db_profile", {})

        # 5) Diccionario de columnas del dataset sobre el DatabaseProfile.
        dict_md_path = None
        dict_json_path = None
        column_dictionary = None
        dict_res = build_column_dictionary(db_profile)
        if dict_res.get("status") == "ok":
            ts = datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S")
            dict_md_path = os.path.join(
                report_dir, f"eda_bq_dataset_{dataset}_{ts}_dict.md"
            )
            dict_json_path = os.path.join(
                report_dir, f"eda_bq_dataset_{dataset}_{ts}_dict.json"
            )
            with open(dict_md_path, "w", encoding="utf-8") as fh:
                fh.write(dict_res.get("markdown", ""))
            # JSON sidecar del diccionario sin el markdown (compacto).
            dict_payload = {k: v for k, v in dict_res.items() if k != "markdown"}
            with open(dict_json_path, "w", encoding="utf-8") as fh:
                fh.write(
                    json.dumps(dict_payload, ensure_ascii=False, indent=1, default=str)
                )
            column_dictionary = dict_payload
        else:
            errors.append(
                {
                    "stage": "column_dictionary",
                    "error": dict_res.get("error", "build_column_dictionary fallo"),
                }
            )

        # 6) Limpieza del DuckDB compartido salvo que se pida conservarlo.
        final_duckdb = duckdb_path
        if not keep_duckdb and os.path.exists(duckdb_path):
            try:
                os.remove(duckdb_path)
            except OSError:
                pass
            final_duckdb = None

        return {
            "status": "ok",
            "project_id": project_id,
            "dataset": dataset,
            "n_tables_loaded": len(loaded_tables),
            "n_tables_profiled": db_profile.get("n_tables", 0),
            "tables_skipped": tables_skipped,
            "errors": errors,
            "duckdb_path": final_duckdb,
            "db_profile": db_profile,
            "column_dictionary": column_dictionary,
            "report_md_path": prof.get("report_md_path"),
            "report_json_path": prof.get("report_json_path"),
            "report_pdf_path": prof.get("report_pdf_path"),
            "dict_md_path": dict_md_path,
            "dict_json_path": dict_json_path,
        }
    except Exception as e:  # noqa: BLE001
        return {"status": "error", "error": str(e), "stage": "unexpected"}