conseguir_datos_con_agentes/notebooks/postgres_con_extensiones.py

import marimo

__generated_with = "0.16.5"
app = marimo.App(width="columns")


@app.cell
def _():
    import marimo as mo
    return (mo,)


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""# 🐘 Postgres: DDBB con extensiones avanzadas (pgvector, PostGIS, TimescaleDB, etc.)""")
    return


@app.cell
def _():
    import os
    import textwrap
    import yaml
    import subprocess
    from pathlib import Path

    # ==== Parámetros ====
    PUERTO_POSTGRES = 55432
    SERVICIO = "postgres_ext"
    PASSWORD = "mipassword"
    USER = "postgres"
    DBNAME = "basededatos"

    # ==== Extensiones que quieres habilitar ====
    EXTENSIONES = [
        # Builtins
        "hstore", "citext", "uuid-ossp", "pg_trgm",
        "fuzzystrmatch", "tablefunc", "unaccent", "ltree",
        # Ecosistema extendido
        "postgis", "pgvector", "timescaledb",
    ]

    RUTA_PROYECTO = Path(".").resolve()
    return (
        DBNAME,
        EXTENSIONES,
        PASSWORD,
        PUERTO_POSTGRES,
        Path,
        RUTA_PROYECTO,
        SERVICIO,
        USER,
        subprocess,
        yaml,
    )


@app.function
def pkgs_para_extensiones(exts, pg_major=15):
    """
    Devuelve (pkgs_apt, builtins, needs_timescale_repo)
    builtins = extensiones que no requieren apt
    pkgs_apt = paquetes apt necesarios para extensiones adicionales
    needs_timescale_repo = True si hay que añadir el repo de TimescaleDB
    """
    builtins = []
    pkgs_apt = []
    needs_timescale_repo = False

    for e in exts:
        e_low = e.lower()

        if e_low in {
            "hstore", "citext", "uuid-ossp", "pg_trgm",
            "fuzzystrmatch", "tablefunc", "unaccent", "ltree"
        }:
            builtins.append(e_low)

        elif e_low == "postgis":
            pkgs_apt += [
                f"postgresql-{pg_major}-postgis-3",
                f"postgresql-{pg_major}-postgis-3-scripts",
                "gdal-bin",
                "proj-bin",
            ]

        elif e_low == "pgvector":
            pkgs_apt += [f"postgresql-{pg_major}-pgvector"]

        elif e_low == "timescaledb":
            needs_timescale_repo = True
            pkgs_apt += [f"timescaledb-2-postgresql-{pg_major}"]

        else:
            raise ValueError(f"Extensión no soportada: {e}")

    pkgs_apt = sorted(set(pkgs_apt))
    return pkgs_apt, builtins, needs_timescale_repo


@app.cell
def _(DBNAME, Path, USER):
    from textwrap import dedent
    from textwrap import indent
    from textwrap import dedent

    def generar_dockerfile(ruta: Path, exts, pg_major=15):
        ruta.mkdir(parents=True, exist_ok=True)
        pkgs_apt, builtins, needs_timescale_repo = pkgs_para_extensiones(exts, pg_major)

        base_image = f"postgres:{pg_major}"

        apt_lines = [
            "RUN apt-get update && \\",
            "    apt-get install -y wget gnupg && \\",
        ]

        if needs_timescale_repo:
            apt_lines += [
                "    wget -qO- https://packagecloud.io/timescale/timescaledb/gpgkey | gpg --dearmor -o /usr/share/keyrings/timescaledb.gpg && \\",
                '    echo \"deb [signed-by=/usr/share/keyrings/timescaledb.gpg] https://packagecloud.io/timescale/timescaledb/debian bookworm main\" > /etc/apt/sources.list.d/timescaledb.list && \\',
                "    apt-get update && \\",
            ]

        if pkgs_apt:
            apt_lines.append("    DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \\")
            for pkg in pkgs_apt:
                apt_lines.append(f"        {pkg} \\")
            apt_lines.append("    && rm -rf /var/lib/apt/lists/*")

        apt_block = "\n".join(apt_lines)

        # Si usa Timescale, añadimos pre-carga
        preload_block = ""
        if "timescaledb" in [e.lower() for e in exts]:
            preload_block = (
                "\n# Preload TimescaleDB\n"
                "RUN echo \"shared_preload_libraries = 'timescaledb'\" >> /usr/share/postgresql/postgresql.conf.sample\n"
            )

        dockerfile = f"""
        FROM {base_image}

        # Variables de entorno
        ENV POSTGRES_USER={USER} \\
            POSTGRES_DB={DBNAME}

        {apt_block}

        {preload_block}

        # Copiamos scripts de inicialización
        COPY docker-entrypoint-initdb.d/ /docker-entrypoint-initdb.d/
        """

        (ruta / "Dockerfile").write_text(dedent(dockerfile).strip() + "\n", encoding="utf-8")

    return (generar_dockerfile,)


@app.cell
def _(Path):
    def generar_init_sql(ruta: Path, exts):
        init_dir = ruta / "docker-entrypoint-initdb.d"
        init_dir.mkdir(parents=True, exist_ok=True)

        lines = ["-- Habilitar extensiones solicitadas"]
        for e in exts:
            e_low = e.lower()
            ext_name = {
                "pgvector": "vector",
                "postgis": "postgis",
                "hstore": "hstore",
                "citext": "citext",
                "uuid-ossp": "\"uuid-ossp\"",
                "pg_trgm": "pg_trgm",
                "fuzzystrmatch": "fuzzystrmatch",
                "tablefunc": "tablefunc",
                "unaccent": "unaccent",
                "ltree": "ltree",
                "timescaledb": "timescaledb",
            }.get(e_low, e_low)
            lines.append(f"CREATE EXTENSION IF NOT EXISTS {ext_name};")

        sql = "\n".join(lines) + "\n"
        (init_dir / "10-extensions.sql").write_text(sql, encoding="utf-8")

    return (generar_init_sql,)


@app.cell
def _(DBNAME, Path, USER, yaml):
    def crear_docker_compose(ruta: Path, servicio: str, puerto_host: int, password: str):
        compose = {
            "services": {
                servicio: {
                    "build": {"context": "."},
                    "restart": "always",
                    "ports": [f"{puerto_host}:5432"],
                    "environment": {
                        "POSTGRES_PASSWORD": password,
                        "POSTGRES_USER": USER,
                        "POSTGRES_DB": DBNAME,
                    },
                    "healthcheck": {
                        "test": ["CMD-SHELL", "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}"],
                        "interval": "10s",
                        "timeout": "5s",
                        "retries": 5,
                    },
                    "volumes": ["postgres_data:/var/lib/postgresql/data"],
                }
            },
            "volumes": {"postgres_data": {}},
        }

        (ruta / "docker-compose.yml").write_text(yaml.dump(compose, sort_keys=False), encoding="utf-8")

    return (crear_docker_compose,)


@app.cell
def _(Path, subprocess):
    def construir_y_levantar(ruta: Path):
        def _run(cmd):
            subprocess.run(cmd, cwd=ruta, check=True)

        try:
            _run(["docker", "compose", "build"])
            _run(["docker", "compose", "up", "-d"])
        except Exception:
            _run(["docker-compose", "build"])
            _run(["docker-compose", "up", "-d"])

    return (construir_y_levantar,)


@app.cell
def _(
    EXTENSIONES,
    PASSWORD,
    PUERTO_POSTGRES,
    RUTA_PROYECTO,
    SERVICIO,
    construir_y_levantar,
    crear_docker_compose,
    generar_dockerfile,
    generar_init_sql,
):
    if __name__ == "__main__":
        RUTA_PROYECTO.mkdir(parents=True, exist_ok=True)
        generar_dockerfile(RUTA_PROYECTO, EXTENSIONES)
        generar_init_sql(RUTA_PROYECTO, EXTENSIONES)
        crear_docker_compose(RUTA_PROYECTO, SERVICIO, PUERTO_POSTGRES, PASSWORD)
        construir_y_levantar(RUTA_PROYECTO)
        print(f"✅ Postgres inicializado en http://localhost:{PUERTO_POSTGRES}")
        print("📦 Extensiones instaladas:", ", ".join(EXTENSIONES))
    return


if __name__ == "__main__":
    app.run()