Upload restore.py with huggingface_hub

restore.py

@@ -0,0 +1,90 @@
+#!/usr/bin/env python3
+"""Restore Open WebUI data into the data dir from the backups dataset.
+
+Runs on every boot. If webui.db is already present the volume is intact and
+nothing happens. If it is missing (volume wiped / factory reset), the latest
+full archive is extracted into the data dir, falling back to the legacy
+single-file DB backup written by older versions of backup.sh.
+
+Never fails the boot: a restore error is logged and ignored, so a first boot
+with no backup yet still starts the app.
+"""
+import os
+import shutil
+import sys
+import tarfile
+
+DATA_DIR = os.environ.get("OWUI_DATA_DIR", "/data")
+REPO_ID = os.environ.get("OWUI_BACKUP_REPO", "tkgray/open-webui-backups")
+ARCHIVE_IN_REPO = "backups/open-webui-data.tar.gz"
+LEGACY_DB_IN_REPO = "backups/webui.db"
+
+DB_PATH = os.path.join(DATA_DIR, "webui.db")
+
+
+def log(msg):
+    print(f"[restore] {msg}", flush=True)
+
+
+def fetch(token, filename):
+    from huggingface_hub import hf_hub_download
+
+    return hf_hub_download(
+        repo_id=REPO_ID,
+        filename=filename,
+        repo_type="dataset",
+        token=token,
+    )
+
+
+def safe_extract(tf, dest):
+    """Extract an archive into dest, refusing absolute paths and .. traversal."""
+    dest = os.path.abspath(dest)
+    for member in tf.getmembers():
+        target = os.path.abspath(os.path.join(dest, member.name))
+        if target != dest and not target.startswith(dest + os.sep):
+            raise ValueError(f"unsafe path in archive: {member.name}")
+    try:
+        tf.extractall(dest, filter="data")  # py3.12+
+    except TypeError:
+        tf.extractall(dest)
+
+
+def main():
+    if os.path.isfile(DB_PATH):
+        log(f"{DB_PATH} already exists — nothing to restore.")
+        return 0
+
+    token = os.environ.get("HF_TOKEN")
+    if not token:
+        log("HF_TOKEN not set — skipping restore.")
+        return 0
+
+    os.makedirs(DATA_DIR, exist_ok=True)
+    log(f"{DB_PATH} missing — attempting restore from {REPO_ID} ...")
+
+    # Preferred: the full archive (DB + vector_db + uploads)
+    try:
+        path = fetch(token, ARCHIVE_IN_REPO)
+        with tarfile.open(path, "r:gz") as tf:
+            safe_extract(tf, DATA_DIR)
+        if os.path.isfile(DB_PATH):
+            log("Full archive restore OK")
+            return 0
+        log("archive extracted but webui.db is missing — falling back")
+    except Exception as exc:  # noqa: BLE001
+        log(f"archive restore failed: {exc}")
+
+    # Fallback: legacy single-file DB backup
+    try:
+        path = fetch(token, LEGACY_DB_IN_REPO)
+        shutil.copy2(path, DB_PATH)
+        log("Legacy DB restore OK")
+    except Exception as exc:  # noqa: BLE001
+        log(f"no backup available yet: {exc}")
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())