From 1a9cb6338360048c9ee493df1e165938da43238d Mon Sep 17 00:00:00 2001
From: Joe Lynch <joe.lynch@aiven.io>
Date: Wed, 13 Nov 2024 14:53:03 +0000
Subject: [PATCH] Add tool for downloading files from a backup.

---
 astacus/manifest.py | 37 +++++++++++++++++++++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/astacus/manifest.py b/astacus/manifest.py
index 53bce957..02ca7346 100644
--- a/astacus/manifest.py
+++ b/astacus/manifest.py
@@ -8,12 +8,18 @@
 
 from astacus.common import ipc
 from astacus.common.rohmustorage import RohmuConfig, RohmuStorage
+from pathlib import Path
 
+import base64
 import json
+import logging
 import msgspec
 import shutil
 import sys
 
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 
 def create_manifest_parsers(parser, subparsers):
     p_manifest = subparsers.add_parser("manifest", help="Examine Astacus backup manifests")
@@ -23,6 +29,7 @@ def create_manifest_parsers(parser, subparsers):
     create_list_parser(manifest_subparsers)
     create_describe_parser(manifest_subparsers)
     create_dump_parser(manifest_subparsers)
+    create_download_files_parser(manifest_subparsers)
 
 
 def create_list_parser(subparsers):
@@ -48,6 +55,36 @@ def create_dump_parser(subparsers):
     p_dump.set_defaults(func=_run_dump)
 
 
+def create_download_files_parser(subparsers):
+    p_download_files = subparsers.add_parser("download-files", help="Download files from a backup manifest")
+    p_download_files.add_argument(
+        "manifest", type=str, help="Manifest object name (can be obtained by running manifest list)"
+    )
+    p_download_files.add_argument("destination", type=str, help="Destination directory to download files to")
+    p_download_files.add_argument("--prefix", type=str, help="Prefix to filter files", required=True)
+    p_download_files.set_defaults(func=_run_download_files)
+
+
+def _run_download_files(args):
+    rohmu_storage = _create_rohmu_storage(args.config, args.storage)
+    manifest = rohmu_storage.download_json(args.manifest, ipc.BackupManifest)
+    destination = Path(args.destination)
+    for snapshot_result in manifest.snapshot_results:
+        assert snapshot_result.state
+        for snapshot_file in snapshot_result.state.files:
+            if not snapshot_file.relative_path.startswith(args.prefix):
+                continue
+
+            path = destination / snapshot_file.relative_path
+            path.parent.mkdir(parents=True, exist_ok=True)
+            logger.info("Downloading %s to %s", snapshot_file.relative_path, path)
+            if snapshot_file.hexdigest:
+                rohmu_storage.download_hexdigest_to_path(snapshot_file.hexdigest, path)
+            else:
+                assert snapshot_file.content_b64 is not None
+                path.write_bytes(base64.b64decode(snapshot_file.content_b64))
+
+
 def _run_list(args):
     rohmu_storage = _create_rohmu_storage(args.config, args.storage)
     json_names = rohmu_storage.list_jsons()