imkrish/remote-postgres

1"""Backup / restore for the in-container Postgres.2 3Two durable targets, either or both:4  * BACKUP_DIR        — local dir (durable only if it's on HF persistent /data)5  * HF_BACKUP_REPO    — a private HF Dataset repo (durable even on the free tier)6 7A fresh boot calls `restore()` to pull the latest dump back; a background scheduler8calls `backup()` every BACKUP_INTERVAL_MIN minutes.9"""10import os11import sys12import glob13import gzip14import time15import threading16import subprocess17from datetime import datetime18 19PG_USER = os.environ.get("POSTGRES_USER", "demo")20PG_DB = os.environ.get("POSTGRES_DB", "demo")21PG_PORT = os.environ.get("PGPORT", "5432")22PGHOST = "/tmp"23 24BACKUP_DIR = os.environ.get("BACKUP_DIR", "/home/appuser/backups")25KEEP = int(os.environ.get("BACKUP_KEEP", "24"))26INTERVAL_MIN = int(os.environ.get("BACKUP_INTERVAL_MIN", "30"))27 28HF_REPO = os.environ.get("HF_BACKUP_REPO", "").strip()      # e.g. "user/pg-backups"29HF_TOKEN = os.environ.get("HF_TOKEN", "").strip()30HF_LATEST = "backups/latest.sql.gz"31 32# Surfaced on the dashboard.33STATE = {"count": 0, "last": None, "last_name": None, "error": None,34         "dir": BACKUP_DIR, "hf": bool(HF_REPO and HF_TOKEN)}35 36 37def _ts():38    return datetime.now().strftime("%Y%m%d-%H%M%S")39 40 41def _hf_api():42    from huggingface_hub import HfApi43    return HfApi(token=HF_TOKEN)44 45 46def ensure_repo():47    if not (HF_REPO and HF_TOKEN):48        return49    try:50        _hf_api().create_repo(repo_id=HF_REPO, repo_type="dataset",51                              private=True, exist_ok=True)52        print(f"[backup] HF dataset ready: {HF_REPO}", flush=True)53    except Exception as e:54        print(f"[backup] create_repo failed: {e}", flush=True)55 56 57def backup():58    """Dump the DB to a gzip file locally, rotate, and mirror to HF (latest)."""59    os.makedirs(BACKUP_DIR, exist_ok=True)60    name = f"backup-{_ts()}.sql.gz"61    path = os.path.join(BACKUP_DIR, name)62    try:63        # whole-cluster dump: every database + roles (not just one DB)64        dump = subprocess.run(65            ["pg_dumpall", "-h", PGHOST, "-p", PG_PORT, "-U", PG_USER],66            stdout=subprocess.PIPE, check=True,67        ).stdout68        with gzip.open(path, "wb") as f:69            f.write(dump)70 71        # keep only the newest KEEP local dumps72        for old in sorted(glob.glob(os.path.join(BACKUP_DIR, "backup-*.sql.gz")))[:-KEEP]:73            try:74                os.remove(old)75            except OSError:76                pass77 78        if HF_REPO and HF_TOKEN:79            try:80                _hf_api().upload_file(81                    path_or_fileobj=path, path_in_repo=HF_LATEST,82                    repo_id=HF_REPO, repo_type="dataset",83                    commit_message=f"backup {name}",84                )85            except Exception as e:86                print(f"[backup] HF upload failed: {e}", flush=True)87 88        STATE.update(count=STATE["count"] + 1, last=datetime.now(),89                     last_name=name, error=None)90        print(f"[backup] wrote {name} ({os.path.getsize(path)} bytes)", flush=True)91        return path92    except Exception as e:93        STATE["error"] = str(e)94        print(f"[backup] FAILED: {e}", flush=True)95        raise96 97 98def restore():99    """Restore the most recent dump (prefers HF, falls back to local). Returns bool."""100    src = None101    if HF_REPO and HF_TOKEN:102        try:103            from huggingface_hub import hf_hub_download104            src = hf_hub_download(HF_REPO, HF_LATEST, repo_type="dataset",105                                  token=HF_TOKEN)106            print(f"[backup] fetched latest dump from HF dataset {HF_REPO}", flush=True)107        except Exception as e:108            print(f"[backup] no HF backup to restore ({e})", flush=True)109    if src is None:110        local = sorted(glob.glob(os.path.join(BACKUP_DIR, "backup-*.sql.gz")))111        src = local[-1] if local else None112    if not src:113        print("[backup] no backup found — starting empty", flush=True)114        return False115 116    print(f"[backup] restoring from {src}", flush=True)117    # connect to 'postgres'; the dumpall script \connects into each database itself118    with gzip.open(src, "rb") as f:119        subprocess.run(120            ["psql", "-h", PGHOST, "-p", PG_PORT, "-U", PG_USER,121             "-d", "postgres", "-v", "ON_ERROR_STOP=0", "-q"],122            input=f.read(), check=True,123        )124    print("[backup] restore complete", flush=True)125    return True126 127 128def _loop():129    while True:130        time.sleep(INTERVAL_MIN * 60)131        try:132            backup()133        except Exception:134            pass135 136 137def start_scheduler():138    """Called from the web app: ensure the HF repo exists and start periodic backups."""139    ensure_repo()140    threading.Thread(target=_loop, daemon=True).start()141    print(f"[backup] scheduler started — every {INTERVAL_MIN} min, "142          f"dir={BACKUP_DIR}, hf={'on' if STATE['hf'] else 'off'}", flush=True)143 144 145if __name__ == "__main__":146    cmd = sys.argv[1] if len(sys.argv) > 1 else ""147    if cmd == "backup":148        backup()149    elif cmd == "restore":150        restore()151    elif cmd == "ensure-repo":152        ensure_repo()153    else:154        print("usage: backup.py [backup|restore|ensure-repo]")155