Spaces:
Running
Running
| #!/usr/bin/env python3 | |
| """Offline-consistent BrainAPI profile backup, verification, and restore.""" | |
| from __future__ import annotations | |
| import argparse | |
| import hashlib | |
| import json | |
| import os | |
| import shlex | |
| import stat | |
| import subprocess | |
| import sys | |
| import tempfile | |
| from datetime import datetime, timezone | |
| from pathlib import Path | |
| from typing import BinaryIO | |
| ROOT = Path(__file__).resolve().parents[1] | |
| DEPLOY = ROOT / "deploy" | |
| ARCHIVER_IMAGE = "alpine:3.21.3" | |
| APP_SERVICES = ("nginx", "brainapi", "brainapi-mcp", "brainapi-worker") | |
| PROFILE_VOLUMES = { | |
| "light": ("redis-data", "brainapi-plugins"), | |
| "heavy": ( | |
| "redis-data", | |
| "milvus-data", | |
| "etcd-data", | |
| "minio-data", | |
| "brainapi-plugins", | |
| ), | |
| } | |
| PROFILE_ALL_VOLUMES = { | |
| "light": ("redis-data", "postgres-data", "brainapi-plugins"), | |
| "heavy": ( | |
| "redis-data", | |
| "neo4j-data", | |
| "neo4j-logs", | |
| "neo4j-plugins", | |
| "etcd-data", | |
| "minio-data", | |
| "milvus-data", | |
| "mongo-data", | |
| "mongo-config", | |
| "brainapi-plugins", | |
| ), | |
| } | |
| RESTORE_ORDER = { | |
| "light": ("redis-data", "postgres", "brainapi-plugins", "stack"), | |
| "heavy": ( | |
| "redis-data", | |
| "etcd-data", | |
| "minio-data", | |
| "milvus-data", | |
| "neo4j", | |
| "mongo", | |
| "brainapi-plugins", | |
| "stack", | |
| ), | |
| } | |
| def _load_env(path: Path) -> dict[str, str]: | |
| values = dict(os.environ) | |
| for raw in path.read_text(encoding="utf-8").splitlines(): | |
| line = raw.strip() | |
| if not line or line.startswith("#") or "=" not in line: | |
| continue | |
| key, value = line.split("=", 1) | |
| values.setdefault(key.strip(), value.strip().strip("'\"")) | |
| return values | |
| class Stack: | |
| def __init__( | |
| self, | |
| profile: str, | |
| env_file: Path, | |
| project: str, | |
| dry_run: bool, | |
| compose_files: list[Path] | None = None, | |
| ): | |
| self.profile = profile | |
| self.env_file = env_file.resolve() | |
| self.project = project | |
| self.dry_run = dry_run | |
| self.compose_files = [DEPLOY / f"docker-compose.{profile}.yaml"] | |
| self.compose_files.extend(path.resolve() for path in (compose_files or [])) | |
| self.env = _load_env(self.env_file) | |
| self.env["BRAINAPI_ENV_FILE"] = str(self.env_file) | |
| self.base = [ | |
| "docker", | |
| "compose", | |
| "--project-name", | |
| project, | |
| "--env-file", | |
| str(self.env_file), | |
| ] | |
| for compose_file in self.compose_files: | |
| self.base.extend(["-f", str(compose_file)]) | |
| def run( | |
| self, | |
| args: list[str], | |
| *, | |
| stdout: BinaryIO | int | None = None, | |
| stdin: BinaryIO | int | None = None, | |
| capture: bool = False, | |
| check: bool = True, | |
| ) -> subprocess.CompletedProcess: | |
| command = [*self.base, *args] | |
| if self.dry_run: | |
| print("+", shlex.join(command)) | |
| return subprocess.CompletedProcess(command, 0, stdout="" if capture else None) | |
| return subprocess.run( | |
| command, | |
| env=self.env, | |
| stdin=stdin, | |
| stdout=subprocess.PIPE if capture else stdout, | |
| stderr=subprocess.PIPE if capture else None, | |
| text=capture, | |
| check=check, | |
| ) | |
| def docker( | |
| self, | |
| args: list[str], | |
| *, | |
| stdout: BinaryIO | int | None = None, | |
| stdin: BinaryIO | int | None = None, | |
| capture: bool = False, | |
| check: bool = True, | |
| ) -> subprocess.CompletedProcess: | |
| command = ["docker", *args] | |
| if self.dry_run: | |
| print("+", shlex.join(command)) | |
| return subprocess.CompletedProcess(command, 0, stdout="" if capture else None) | |
| return subprocess.run( | |
| command, | |
| env=self.env, | |
| stdin=stdin, | |
| stdout=subprocess.PIPE if capture else stdout, | |
| stderr=subprocess.PIPE if capture else None, | |
| text=capture, | |
| check=check, | |
| ) | |
| def volume_name(self, key: str) -> str: | |
| return f"{self.project}_{key}" | |
| def _sha256(path: Path) -> str: | |
| digest = hashlib.sha256() | |
| with path.open("rb") as stream: | |
| for chunk in iter(lambda: stream.read(1024 * 1024), b""): | |
| digest.update(chunk) | |
| return digest.hexdigest() | |
| def _write_stream(stack: Stack, args: list[str], destination: Path) -> None: | |
| if stack.dry_run: | |
| stack.run(args) | |
| destination.touch() | |
| return | |
| with destination.open("wb") as output: | |
| stack.run(args, stdout=output) | |
| def _dump_postgres_brain_databases(stack: Stack, destination: Path) -> None: | |
| script = r''' | |
| set -eu | |
| tmp="$(mktemp -d)" | |
| trap 'rm -rf "$tmp"' EXIT | |
| psql -U "$POSTGRES_USER" -d postgres -Atc \ | |
| "SELECT datname FROM pg_database WHERE datistemplate = false AND left(datname, 6) = 'brain_' ORDER BY datname" | | |
| while IFS= read -r database; do | |
| test -n "$database" || continue | |
| pg_dump -U "$POSTGRES_USER" -d "$database" --format=custom --file="$tmp/$database.dump" | |
| done | |
| tar -C "$tmp" -czf - . | |
| '''.strip() | |
| _write_stream( | |
| stack, | |
| ["exec", "-T", "postgres", "sh", "-c", script], | |
| destination, | |
| ) | |
| def _restore_postgres_brain_databases( | |
| stack: Stack, archive: Path, pg_user: str | |
| ) -> None: | |
| with tempfile.TemporaryDirectory(prefix="brainapi-postgres-restore-") as tmp: | |
| tmp_path = Path(tmp) | |
| if not stack.dry_run: | |
| subprocess.run( | |
| ["tar", "-C", str(tmp_path), "-xzf", str(archive)], check=True | |
| ) | |
| dumps = sorted(tmp_path.glob("brain_*.dump")) | |
| if stack.dry_run: | |
| dumps = [tmp_path / "brain_example.dump"] | |
| for dump in dumps: | |
| database = dump.stem | |
| if not database.replace("_", "").isalnum() or not database.startswith( | |
| "brain_" | |
| ): | |
| raise RuntimeError(f"Invalid PostgreSQL brain dump name: {dump.name}") | |
| stack.run( | |
| ["exec", "-T", "postgres", "createdb", "-U", pg_user, database], | |
| check=False, | |
| ) | |
| if stack.dry_run: | |
| stack.run( | |
| [ | |
| "exec", | |
| "-T", | |
| "postgres", | |
| "pg_restore", | |
| "-U", | |
| pg_user, | |
| "-d", | |
| database, | |
| "--clean", | |
| "--if-exists", | |
| ] | |
| ) | |
| else: | |
| with dump.open("rb") as source: | |
| stack.run( | |
| [ | |
| "exec", | |
| "-T", | |
| "postgres", | |
| "pg_restore", | |
| "-U", | |
| pg_user, | |
| "-d", | |
| database, | |
| "--clean", | |
| "--if-exists", | |
| ], | |
| stdin=source, | |
| ) | |
| def _archive_volume(stack: Stack, key: str, destination: Path) -> None: | |
| if stack.dry_run: | |
| stack.docker(["run", "--rm", "-v", f"{stack.volume_name(key)}:/source:ro", ARCHIVER_IMAGE, "tar", "-C", "/source", "-czf", "-", "."]) | |
| destination.touch() | |
| return | |
| with destination.open("wb") as output: | |
| stack.docker( | |
| [ | |
| "run", | |
| "--rm", | |
| "-v", | |
| f"{stack.volume_name(key)}:/source:ro", | |
| ARCHIVER_IMAGE, | |
| "tar", | |
| "-C", | |
| "/source", | |
| "-czf", | |
| "-", | |
| ".", | |
| ], | |
| stdout=output, | |
| ) | |
| def _restore_volume(stack: Stack, key: str, archive: Path) -> None: | |
| if stack.dry_run: | |
| stack.docker(["run", "--rm", "-i", "-v", f"{stack.volume_name(key)}:/target", ARCHIVER_IMAGE, "tar", "-C", "/target", "-xzf", "-"]) | |
| return | |
| with archive.open("rb") as source: | |
| stack.docker( | |
| [ | |
| "run", | |
| "--rm", | |
| "-i", | |
| "-v", | |
| f"{stack.volume_name(key)}:/target", | |
| ARCHIVER_IMAGE, | |
| "tar", | |
| "-C", | |
| "/target", | |
| "-xzf", | |
| "-", | |
| ], | |
| stdin=source, | |
| ) | |
| def _running_services(stack: Stack) -> list[str]: | |
| result = stack.run( | |
| ["ps", "--services", "--filter", "status=running"], capture=True | |
| ) | |
| return [line for line in (result.stdout or "").splitlines() if line] | |
| def _current_image(stack: Stack) -> str: | |
| result = stack.run(["config", "--format", "json"], capture=True) | |
| if stack.dry_run: | |
| return stack.env.get("BRAINAPI_IMAGE", "dry-run") | |
| config = json.loads(result.stdout) | |
| return str(config["services"]["brainapi"]["image"]) | |
| def _manifest(backup_dir: Path) -> dict: | |
| return json.loads((backup_dir / "manifest.json").read_text(encoding="utf-8")) | |
| def verify_backup(backup_dir: Path, expected_profile: str | None = None) -> dict: | |
| manifest = _manifest(backup_dir) | |
| if manifest.get("format_version") != 1: | |
| raise RuntimeError("Unsupported backup manifest format") | |
| if expected_profile and manifest.get("profile") != expected_profile: | |
| raise RuntimeError("Backup profile does not match the requested profile") | |
| for name, expected in manifest.get("checksums", {}).items(): | |
| path = backup_dir / name | |
| if not path.is_file(): | |
| raise RuntimeError(f"Missing backup artifact: {name}") | |
| actual = _sha256(path) | |
| if actual != expected: | |
| raise RuntimeError(f"Checksum mismatch: {name}") | |
| return manifest | |
| def backup(stack: Stack, backup_root: Path) -> Path: | |
| timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") | |
| backup_root = backup_root.expanduser().resolve() | |
| backup_root.mkdir(parents=True, mode=0o700, exist_ok=True) | |
| backup_root.chmod(stat.S_IRWXU) | |
| destination = backup_root / f"{stack.profile}-{timestamp}" | |
| destination.mkdir(parents=True, mode=0o700, exist_ok=False) | |
| destination.chmod(stat.S_IRWXU) | |
| running = _running_services(stack) | |
| try: | |
| stack.run(["stop", *APP_SERVICES]) | |
| if stack.profile == "light": | |
| pg_user = stack.env.get("POSTGRES_USERNAME", "brainapi") | |
| pg_db = stack.env.get("POSTGRES_SYSTEM_DATABASE", "brainapi") | |
| _write_stream( | |
| stack, | |
| ["exec", "-T", "postgres", "pg_dump", "-U", pg_user, "-d", pg_db, "--format=custom"], | |
| destination / "postgres.dump", | |
| ) | |
| _dump_postgres_brain_databases( | |
| stack, destination / "postgres-brains.tar.gz" | |
| ) | |
| stack.run(["exec", "-T", "redis", "sh", "-c", 'redis-cli -a "$REDIS_PASSWORD" SAVE']) | |
| stack.run(["stop", "redis", "postgres"]) | |
| else: | |
| mongo_user = stack.env.get("MONGO_USERNAME", "brainapi") | |
| _write_stream( | |
| stack, | |
| ["exec", "-T", "mongo", "sh", "-c", f'mongodump --archive --gzip --username {shlex.quote(mongo_user)} --password "$MONGO_INITDB_ROOT_PASSWORD" --authenticationDatabase admin'], | |
| destination / "mongo.archive.gz", | |
| ) | |
| stack.run(["exec", "-T", "redis", "sh", "-c", 'redis-cli -a "$REDIS_PASSWORD" SAVE']) | |
| stack.run(["stop", "neo4j", "mongo", "milvus", "etcd", "minio", "redis"]) | |
| neo_dir = destination / "neo4j" | |
| neo_dir.mkdir(mode=0o700) | |
| stack.run( | |
| ["run", "--rm", "--no-deps", "-v", f"{neo_dir}:/backup", "neo4j", "neo4j-admin", "database", "dump", "system", "neo4j", "--to-path=/backup", "--overwrite-destination=true"] | |
| ) | |
| if not stack.dry_run: | |
| subprocess.run(["tar", "-C", str(neo_dir), "-czf", str(destination / "neo4j.dumps.tar.gz"), "."], check=True) | |
| for child in neo_dir.iterdir(): | |
| child.unlink() | |
| neo_dir.rmdir() | |
| for volume in PROFILE_VOLUMES[stack.profile]: | |
| _archive_volume(stack, volume, destination / f"{volume}.tar.gz") | |
| artifacts = sorted(path for path in destination.iterdir() if path.is_file()) | |
| manifest = { | |
| "format_version": 1, | |
| "profile": stack.profile, | |
| "brainapi_image": _current_image(stack), | |
| "brainapi_version": _current_image(stack).rsplit(":", 1)[-1], | |
| "timestamp": timestamp, | |
| "services": ( | |
| ["postgres", "redis", "brainapi-plugins"] | |
| if stack.profile == "light" | |
| else ["neo4j", "mongo", "redis", "milvus", "etcd", "minio", "brainapi-plugins"] | |
| ), | |
| "restore_order": list(RESTORE_ORDER[stack.profile]), | |
| "checksums": {path.name: _sha256(path) for path in artifacts}, | |
| } | |
| (destination / "manifest.json").write_text( | |
| json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8" | |
| ) | |
| verify_backup(destination, stack.profile) | |
| return destination | |
| finally: | |
| if running: | |
| stack.run(["start", *running], check=False) | |
| def _assert_empty_stopped(stack: Stack) -> None: | |
| running = _running_services(stack) | |
| if running: | |
| raise RuntimeError("Restore target must be stopped; running: " + ", ".join(running)) | |
| for key in PROFILE_ALL_VOLUMES[stack.profile]: | |
| volume = stack.volume_name(key) | |
| exists = stack.docker(["volume", "inspect", volume], check=False).returncode == 0 | |
| if not exists or stack.dry_run: | |
| continue | |
| result = stack.docker( | |
| ["run", "--rm", "-v", f"{volume}:/source:ro", ARCHIVER_IMAGE, "sh", "-c", "test -z \"$(find /source -mindepth 1 -maxdepth 1 -print -quit)\""], | |
| check=False, | |
| ) | |
| if result.returncode != 0: | |
| raise RuntimeError(f"Restore target volume is not empty: {volume}") | |
| def restore(stack: Stack, backup_dir: Path) -> None: | |
| manifest = verify_backup(backup_dir, stack.profile) | |
| if manifest["brainapi_image"] != _current_image(stack): | |
| raise RuntimeError("Backup image does not match BRAINAPI_IMAGE") | |
| _assert_empty_stopped(stack) | |
| stack.run(["up", "--no-start"]) | |
| for key in PROFILE_VOLUMES[stack.profile]: | |
| _restore_volume(stack, key, backup_dir / f"{key}.tar.gz") | |
| if stack.profile == "light": | |
| stack.run(["up", "-d", "--wait", "postgres"]) | |
| pg_user = stack.env.get("POSTGRES_USERNAME", "brainapi") | |
| pg_db = stack.env.get("POSTGRES_SYSTEM_DATABASE", "brainapi") | |
| if stack.dry_run: | |
| stack.run(["exec", "-T", "postgres", "pg_restore", "-U", pg_user, "-d", pg_db, "--clean", "--if-exists"]) | |
| else: | |
| with (backup_dir / "postgres.dump").open("rb") as source: | |
| stack.run(["exec", "-T", "postgres", "pg_restore", "-U", pg_user, "-d", pg_db, "--clean", "--if-exists"], stdin=source) | |
| _restore_postgres_brain_databases( | |
| stack, backup_dir / "postgres-brains.tar.gz", pg_user | |
| ) | |
| else: | |
| with tempfile.TemporaryDirectory(prefix="brainapi-neo4j-restore-") as tmp: | |
| tmp_path = Path(tmp) | |
| if not stack.dry_run: | |
| subprocess.run(["tar", "-C", str(tmp_path), "-xzf", str(backup_dir / "neo4j.dumps.tar.gz")], check=True) | |
| stack.run(["run", "--rm", "--no-deps", "-v", f"{tmp_path}:/backup:ro", "neo4j", "neo4j-admin", "database", "load", "system", "neo4j", "--from-path=/backup", "--overwrite-destination=true"]) | |
| stack.run(["up", "-d", "--wait", "mongo"]) | |
| mongo_user = stack.env.get("MONGO_USERNAME", "brainapi") | |
| if stack.dry_run: | |
| stack.run(["exec", "-T", "mongo", "sh", "-c", f'mongorestore --archive --gzip --drop --username {shlex.quote(mongo_user)} --password "$MONGO_INITDB_ROOT_PASSWORD" --authenticationDatabase admin']) | |
| else: | |
| with (backup_dir / "mongo.archive.gz").open("rb") as source: | |
| stack.run(["exec", "-T", "mongo", "sh", "-c", f'mongorestore --archive --gzip --drop --username {shlex.quote(mongo_user)} --password "$MONGO_INITDB_ROOT_PASSWORD" --authenticationDatabase admin'], stdin=source) | |
| stack.run(["up", "-d", "--wait"]) | |
| def _parser() -> argparse.ArgumentParser: | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument("command", choices=("backup", "verify", "restore")) | |
| parser.add_argument("--profile", choices=("light", "heavy"), required=True) | |
| parser.add_argument("--env-file", type=Path, required=True) | |
| parser.add_argument("--project-name", default="brainapi") | |
| parser.add_argument("--backup-dir", type=Path, default=Path(os.getenv("BACKUP_DIR", "/srv/brainapi/backups"))) | |
| parser.add_argument("--archive", type=Path, help="Backup directory for verify/restore") | |
| parser.add_argument( | |
| "--compose-file", | |
| action="append", | |
| default=[], | |
| type=Path, | |
| help="Additional Compose override file; may be repeated.", | |
| ) | |
| parser.add_argument("--dry-run", action="store_true") | |
| return parser | |
| def main() -> int: | |
| args = _parser().parse_args() | |
| try: | |
| if args.command == "verify": | |
| if not args.archive: | |
| raise RuntimeError("--archive is required for verify") | |
| verify_backup(args.archive.resolve(), args.profile) | |
| print(f"Verified {args.archive}") | |
| return 0 | |
| stack = Stack( | |
| args.profile, | |
| args.env_file, | |
| args.project_name, | |
| args.dry_run, | |
| args.compose_file, | |
| ) | |
| if args.command == "backup": | |
| destination = backup(stack, args.backup_dir) | |
| print(destination) | |
| else: | |
| if not args.archive: | |
| raise RuntimeError("--archive is required for restore") | |
| restore(stack, args.archive.resolve()) | |
| print(f"Restored {args.archive}") | |
| return 0 | |
| except (KeyError, OSError, RuntimeError, subprocess.CalledProcessError) as exc: | |
| print(f"backup_restore: {exc}", file=sys.stderr) | |
| return 1 | |
| if __name__ == "__main__": | |
| raise SystemExit(main()) | |