myrag/Generacion de bbdd.py

import marimo

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


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

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

    # Lista de extensiones que quieres habilitar
    EXTENSIONES = [
        # builtin
        "hstore", "citext", "uuid-ossp", "pg_trgm",
        # requieren paquetes
        # "postgis",
        "pgvector",
        # "timescaledb"  # <- si quieres usar base de timescaledb, activa la bandera abajo
    ]

    # Usa imagen base de timescaledb cuando la extensión 'timescaledb' esté en la lista
    timescaledb_base_image = False  # pon True si quieres usar la imagen base de TimescaleDB

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


@app.function
def pkgs_para_extensiones(exts, pg_major=15, use_timescale_base=False):
    """
    Devuelve (pkgs_apt, builtins) para las extensiones solicitadas.
    builtins = extensiones que no requieren apt
    pkgs_apt = paquetes apt necesarios para otras extensiones
    """
    builtins = []
    pkgs_apt = []

    for e in exts:
        e_low = e.lower()
        if e_low in {"hstore", "citext", "uuid-ossp", "pg_trgm"}:
            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":
            if not use_timescale_base:
                # Para instalar timescaledb necesitarías repos adicionales
                pass
        else:
            raise ValueError(f"Extensión no soportada en este helper: {e}")

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


@app.cell
def _(DBNAME, Path, USER, textwrap):
    def generar_dockerfile(ruta: Path, exts, use_timescale_base=False, pg_major=15):
        ruta.mkdir(parents=True, exist_ok=True)
        pkgs_apt, builtins = pkgs_para_extensiones(exts, pg_major, use_timescale_base)

        base_image = (
            f"timescale/timescaledb:2.16-pg{pg_major}" if use_timescale_base else f"postgres:{pg_major}"
        )

        apt_block = ""
        if pkgs_apt:
            apt_lines = [
                "RUN apt-get update && \\",
                "    DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \\",
            ]
            # 👇 Backslash en **todas** las líneas de paquetes
            for pkg in pkgs_apt:
                apt_lines.append(f"        {pkg} \\")
            # y ya colgamos el rm del último paquete
            apt_lines.append("    && rm -rf /var/lib/apt/lists/*")
            apt_block = "\n".join(apt_lines)

        dockerfile = f"""
        FROM {base_image}

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

        {apt_block}

        # Copiamos scripts de inicialización
        COPY docker-entrypoint-initdb.d/ /docker-entrypoint-initdb.d/
        """
        (ruta / "Dockerfile").write_text(textwrap.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",
                "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 = {
            "version": "3.8",
            "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,
    timescaledb_base_image,
):
    if __name__ == "__main__":
        RUTA_PROYECTO.mkdir(parents=True, exist_ok=True)
        generar_dockerfile(RUTA_PROYECTO, EXTENSIONES, timescaledb_base_image)
        generar_init_sql(RUTA_PROYECTO, EXTENSIONES)
        crear_docker_compose(RUTA_PROYECTO, SERVICIO, PUERTO_POSTGRES, PASSWORD)
        construir_y_levantar(RUTA_PROYECTO)
    return


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