Spaces:
Running on Zero
Running on Zero
| """Run private application backup/restore and retention outside serving requests.""" | |
| from __future__ import annotations | |
| import argparse | |
| import hashlib | |
| import json | |
| import os | |
| from datetime import UTC, datetime | |
| from pathlib import Path | |
| from shared.database import StorageUnavailable | |
| from shared.storage_ops import ( | |
| BackupError, | |
| backup, | |
| cleanup, | |
| key_file, | |
| prune_backups, | |
| read_backup, | |
| restore, | |
| ) | |
| ROOT = Path(__file__).resolve().parents[1] | |
| DEFAULT_KEY = ROOT / "deploy/.generated/precedent-backup.key" | |
| DEFAULT_BACKUPS = ROOT / "deploy/.generated/storage-backups" | |
| def main() -> None: | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument("--key-file", type=Path, default=DEFAULT_KEY) | |
| subcommands = parser.add_subparsers(dest="operation", required=True) | |
| subcommands.add_parser( | |
| "init-key", help="Create a private encryption key; never print or overwrite it." | |
| ) | |
| export = subcommands.add_parser( | |
| "backup", help="Encrypt a consistent application-data snapshot." | |
| ) | |
| export.add_argument("--output", type=Path) | |
| for name in ("inspect", "restore"): | |
| command = subcommands.add_parser(name) | |
| command.add_argument("--backup", type=Path, required=True) | |
| maintenance = subcommands.add_parser("cleanup") | |
| maintenance.add_argument("--batch", type=int, default=500) | |
| pruning = subcommands.add_parser("prune-backups") | |
| pruning.add_argument("--directory", type=Path, default=DEFAULT_BACKUPS) | |
| args = parser.parse_args() | |
| try: | |
| if args.operation == "init-key": | |
| key_file(args.key_file, create=True) | |
| print(json.dumps({"status": "key_created", "path": str(args.key_file)})) | |
| return | |
| if args.operation == "inspect": | |
| document = read_backup(args.backup, key_file(args.key_file)) | |
| report = { | |
| "version": document["version"], | |
| "created_at": document["created_at"], | |
| "rows": {table: len(rows) for table, rows in document["data"].items()}, | |
| } | |
| elif args.operation == "prune-backups": | |
| report = prune_backups(args.directory, key_file(args.key_file)) | |
| else: | |
| dsn = os.environ.get("LEGAL_AI_QUOTA_DATABASE_URL") | |
| if not dsn: | |
| raise BackupError( | |
| "Set LEGAL_AI_QUOTA_DATABASE_URL privately in the operator environment." | |
| ) | |
| if args.operation == "cleanup": | |
| report = cleanup(dsn, batch=args.batch) | |
| elif args.operation == "backup": | |
| destination = args.output or DEFAULT_BACKUPS / datetime.now(UTC).strftime( | |
| "precedent-%Y%m%dT%H%M%S%fZ.pab" | |
| ) | |
| report = backup(dsn, destination, key_file(args.key_file)) | |
| report.update( | |
| path=str(destination), | |
| sha256=hashlib.sha256(destination.read_bytes()).hexdigest(), | |
| ) | |
| else: | |
| report = restore(dsn, args.backup, key_file(args.key_file)) | |
| print(json.dumps({"operation": args.operation, **report})) | |
| except (BackupError, StorageUnavailable) as exc: | |
| raise SystemExit(str(exc)) from None | |
| except (OSError, ImportError): | |
| raise SystemExit( | |
| "Storage maintenance failed. Check private file access and the storage/ops extras." | |
| ) from None | |
| if __name__ == "__main__": | |
| main() | |