Spaces:
Paused
Paused
File size: 2,653 Bytes
4da2f34 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | 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())
|