deploy/huggingface/sqlite_backup.py

from __future__ import annotations

import argparse
import os
import shutil
import sqlite3
import tempfile
from pathlib import Path


def _temp_path_for(destination: Path) -> Path:
    destination.parent.mkdir(parents=True, exist_ok=True)
    fd, temp_name = tempfile.mkstemp(prefix=destination.stem + ".", suffix=".tmp", dir=destination.parent)
    os.close(fd)
    return Path(temp_name)


def backup_database(source: Path | str, destination: Path | str) -> bool:
    source_path = Path(source)
    destination_path = Path(destination)
    if not source_path.exists():
        print(f"[hf-space] backup skipped; source missing: {source_path}")
        return False

    temp_path = _temp_path_for(destination_path)
    src_conn = None
    dst_conn = None
    try:
        src_conn = sqlite3.connect(source_path)
        dst_conn = sqlite3.connect(temp_path)
        src_conn.backup(dst_conn)
        dst_conn.commit()
        dst_conn.close()
        dst_conn = None
        src_conn.close()
        src_conn = None
        os.replace(temp_path, destination_path)
        print(f"[hf-space] backup updated: {destination_path}")
        return True
    finally:
        if dst_conn is not None:
            dst_conn.close()
        if src_conn is not None:
            src_conn.close()
        if temp_path.exists():
            temp_path.unlink(missing_ok=True)


def restore_database(source: Path | str, destination: Path | str) -> bool:
    source_path = Path(source)
    destination_path = Path(destination)
    if not source_path.exists():
        print(f"[hf-space] restore skipped; source missing: {source_path}")
        return False

    temp_path = _temp_path_for(destination_path)
    try:
        shutil.copy2(source_path, temp_path)
        os.replace(temp_path, destination_path)
        print(f"[hf-space] restore completed: {destination_path}")
        return True
    finally:
        if temp_path.exists():
            temp_path.unlink(missing_ok=True)


def main() -> int:
    parser = argparse.ArgumentParser(description="Backup or restore a SQLite database for HF Spaces.")
    subparsers = parser.add_subparsers(dest="command", required=True)

    for command in ("backup", "restore"):
        command_parser = subparsers.add_parser(command)
        command_parser.add_argument("--source", required=True)
        command_parser.add_argument("--destination", required=True)

    args = parser.parse_args()
    if args.command == "backup":
        backup_database(args.source, args.destination)
    else:
        restore_database(args.source, args.destination)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())