imkrish/remote-postgres

1import os2import json3import time4import secrets5import threading6from datetime import datetime, timedelta7 8import psycopg9import gradio as gr10 11import backup12from fastapi import FastAPI, Depends, HTTPException, status13from fastapi.responses import JSONResponse14from fastapi.security import HTTPBasic, HTTPBasicCredentials15 16PG_USER = os.environ.get("POSTGRES_USER", "demo")17PG_PASSWORD = os.environ.get("POSTGRES_PASSWORD", "")18PG_DB = os.environ.get("POSTGRES_DB", "demo")19PG_PORT = os.environ.get("PGPORT", "5432")20 21# If set, the Gradio UI and the JSON API require this password (username: admin).22APP_PASSWORD = os.environ.get("APP_PASSWORD", "")23APP_USER = os.environ.get("APP_USER", "admin")24 25TUNNEL_FILE = "/tmp/tunnel.json"   # written by start.sh (bore or ngrok)26 27BOOT = datetime.now()28HISTORY: list[tuple[str, bool, bool]] = []   # (HH:MM:SS, postgres_up, tunnel_up)29HISTORY_MAX = 50030 31# Keep-alive counter (incremented by the in-container loop AND any external cron).32KEEPALIVE = {"count": 0, "last": None, "src": "—"}33 34 35# ---------------------------------------------------------------- helpers ----36def get_tunnel():37    """Return (host, port) of the public TCP tunnel from /tmp/tunnel.json, or None."""38    try:39        with open(TUNNEL_FILE) as f:40            d = json.load(f)41        if d.get("host") and d.get("port"):42            return d["host"], str(d["port"])43    except Exception:44        return None45    return None46 47 48def pg_ok() -> bool:49    try:50        with psycopg.connect(51            host="/tmp", port=PG_PORT, user=PG_USER,52            dbname=PG_DB, password=PG_PASSWORD, connect_timeout=3,53        ):54            return True55    except Exception:56        return False57 58 59def build_url(tunnel):60    if not tunnel:61        return None62    host, port = tunnel63    return f"postgresql://{PG_USER}:{PG_PASSWORD}@{host}:{port}/{PG_DB}"64 65 66def humanize(delta: timedelta) -> str:67    s = int(delta.total_seconds())68    d, s = divmod(s, 86400)69    h, s = divmod(s, 3600)70    m, s = divmod(s, 60)71    if d:72        return f"{d}d {h}h {m}m"73    if h:74        return f"{h}h {m}m {s}s"75    return f"{m}m {s}s"76 77 78def record_check():79    ok = pg_ok()80    tn = get_tunnel() is not None81    HISTORY.append((datetime.now().strftime("%H:%M:%S"), ok, tn))82    del HISTORY[:-HISTORY_MAX]83    return ok, tn84 85 86def _watcher_loop():87    while True:88        try:89            record_check()90        except Exception:91            pass92        time.sleep(30)93 94 95# Background uptime watcher — records a sample every 30s while the container lives.96threading.Thread(target=_watcher_loop, daemon=True).start()97 98# Periodic Postgres backups (to BACKUP_DIR and/or a private HF Dataset).99backup.start_scheduler()100 101 102# ------------------------------------------------------------- JSON / API ----103api = FastAPI(title="Remote Postgres", docs_url=None, redoc_url=None)104security = HTTPBasic(auto_error=False)105 106 107def require_auth(credentials: HTTPBasicCredentials | None = Depends(security)):108    if not APP_PASSWORD:109        return110    if credentials is None or not (111        secrets.compare_digest(credentials.username, APP_USER)112        and secrets.compare_digest(credentials.password, APP_PASSWORD)113    ):114        raise HTTPException(115            status_code=status.HTTP_401_UNAUTHORIZED,116            detail="Authentication required",117            headers={"WWW-Authenticate": "Basic"},118        )119 120 121@api.get("/health")122async def health():123    """Open endpoint — used by the self keep-alive and external uptime monitors."""124    tunnel = get_tunnel()125    ok = pg_ok()126    return JSONResponse({127        "status": "ok" if ok else "degraded",128        "postgres": ok,129        "tunnel": bool(tunnel),130        "uptime": humanize(datetime.now() - BOOT),131    })132 133 134@api.get("/keepalive")135async def keepalive(src: str = "self"):136    """Hit to reset HF's idle timer (by the in-container loop AND any external cron).137    Open + logged so you can confirm pings are landing."""138    KEEPALIVE["count"] += 1139    KEEPALIVE["last"] = datetime.now()140    KEEPALIVE["src"] = src141    print(f"[keepalive] hit #{KEEPALIVE['count']} from '{src}' at "142          f"{KEEPALIVE['last'].strftime('%Y-%m-%d %H:%M:%S')}", flush=True)143    return {"ok": True, "count": KEEPALIVE["count"], "src": src}144 145 146@api.get("/api/connection")147async def connection(_: None = Depends(require_auth)):148    tunnel = get_tunnel()149    host, port = tunnel if tunnel else (None, None)150    return {151        "postgres_up": pg_ok(),152        "connection_url": build_url(tunnel),153        "host": host, "port": port,154        "user": PG_USER, "password": PG_PASSWORD, "database": PG_DB,155    }156 157 158@api.post("/admin/backup")159async def admin_backup(_: None = Depends(require_auth)):160    try:161        path = backup.backup()162        return {"ok": True, "file": os.path.basename(path)}163    except Exception as e:164        return JSONResponse({"ok": False, "error": str(e)}, status_code=500)165 166 167@api.post("/admin/restore")168async def admin_restore(_: None = Depends(require_auth)):169    """Reload the latest backup (whole cluster). Cleanest on a fresh/empty DB;170    on a populated DB it best-effort merges and skips rows that already exist."""171    try:172        done = backup.restore()173        return {"ok": bool(done)}174    except Exception as e:175        return JSONResponse({"ok": False, "error": str(e)}, status_code=500)176 177 178# ------------------------------------------------- Gradio uptime watcher ----179def snapshot():180    tunnel = get_tunnel()181    up = pg_ok()182    url = build_url(tunnel) or "— tunnel starting… refresh in a few seconds —"183    host, port = tunnel if tunnel else ("—", "—")184 185    samples = len(HISTORY)186    pct = (100 * sum(1 for _, o, _ in HISTORY if o) / samples) if samples else 100.0187    status_md = (188        f"### Status\n"189        f"- **Postgres:** {'🟢 up' if up else '🔴 down'}\n"190        f"- **Public tunnel:** {'🟢 up' if tunnel else '🟡 off'}\n"191        f"- **Container uptime:** {humanize(datetime.now() - BOOT)}\n"192        f"- **PG availability (this container):** {pct:.1f}%  ·  {samples} checks\n"193        f"- **Keep-alive hits:** {KEEPALIVE['count']}  ·  last from "194        f"`{KEEPALIVE['src']}` at "195        f"{KEEPALIVE['last'].strftime('%H:%M:%S') if KEEPALIVE['last'] else 'never'}\n"196        f"- **Backups:** {backup.STATE['count']} taken  ·  last: "197        f"{backup.STATE['last'].strftime('%H:%M:%S') if backup.STATE['last'] else 'none yet'}"198        f"  ·  HF off-Space: {'on' if backup.STATE['hf'] else 'off'}\n"199        f"- **Last checked:** {datetime.now().strftime('%H:%M:%S')}"200    )201    table = [202        [t, "✅" if o else "❌", "✅" if n else "❌"]203        for (t, o, n) in reversed(HISTORY[-25:])204    ]205    return status_md, url, host, port, PG_USER, PG_PASSWORD, PG_DB, table206 207 208with gr.Blocks(title="Remote Postgres — Uptime Watcher", theme=gr.themes.Soft()) as demo:209    gr.Markdown("# 🐘 Remote Postgres — Uptime Watcher")210    gr.Markdown(211        "Throwaway Postgres for demo projects, exposed over a TCP tunnel (bore). "212        "Data is backed up, but the URL **changes on every restart**."213    )214    status_box = gr.Markdown()215    url_box = gr.Textbox(216        label="Connection URL — paste into your projects",217        interactive=False, show_copy_button=True,218    )219    with gr.Row():220        host_box = gr.Textbox(label="Host", interactive=False)221        port_box = gr.Textbox(label="Port", interactive=False)222    with gr.Row():223        user_box = gr.Textbox(label="User", interactive=False)224        pw_box = gr.Textbox(label="Password", interactive=False)225        db_box = gr.Textbox(label="Database", interactive=False)226    history_box = gr.Dataframe(227        headers=["time", "postgres", "tunnel"],228        label="Recent checks (newest first)",229        interactive=False,230    )231    with gr.Row():232        refresh_btn = gr.Button("Refresh now", variant="primary")233        backup_btn = gr.Button("Back up now")234        restore_btn = gr.Button("Restore latest")235 236    def backup_now():237        try:238            backup.backup()239        except Exception as e:240            print(f"[ui] manual backup failed: {e}", flush=True)241        return snapshot()242 243    def restore_now():244        try:245            backup.restore()246        except Exception as e:247            print(f"[ui] manual restore failed: {e}", flush=True)248        return snapshot()249 250    outputs = [status_box, url_box, host_box, port_box,251               user_box, pw_box, db_box, history_box]252    timer = gr.Timer(10)253    timer.tick(snapshot, outputs=outputs)254    refresh_btn.click(snapshot, outputs=outputs)255    backup_btn.click(backup_now, outputs=outputs)256    restore_btn.click(restore_now, outputs=outputs)257    demo.load(snapshot, outputs=outputs)258 259 260auth = (APP_USER, APP_PASSWORD) if APP_PASSWORD else None261app = gr.mount_gradio_app(api, demo, path="/", auth=auth)262