From 5d5ceee1cf68b583ca9961883d44c30dbb4353e3 Mon Sep 17 00:00:00 2001
From: Houkime <>
Date: Wed, 22 Feb 2023 18:48:08 +0000
Subject: [PATCH] feat(backups): sizing up snapshots

---
 selfprivacy_api/backup/__init__.py        | 10 +++++++
 selfprivacy_api/backup/backuper.py        |  4 +++
 selfprivacy_api/backup/restic_backuper.py | 35 ++++++++++++++++++++---
 tests/test_graphql/test_backup.py         |  8 ++++++
 4 files changed, 53 insertions(+), 4 deletions(-)

diff --git a/selfprivacy_api/backup/__init__.py b/selfprivacy_api/backup/__init__.py
index 4261e35..903e38b 100644
--- a/selfprivacy_api/backup/__init__.py
+++ b/selfprivacy_api/backup/__init__.py
@@ -77,3 +77,13 @@ class Backups:
         self.restore_service_from_snapshot(
             get_service_by_id(snapshot.service_name), snapshot.id
         )
+
+    def service_snapshot_size(self, service: Service, snapshot_id: str) -> float:
+        repo_name = service.get_id()
+        return self.provider.backuper.restored_size(repo_name, snapshot_id)
+
+    # Our dummy service is not yet globally registered so this is not testable yet
+    def snapshot_restored_size(self, snapshot: Snapshot) -> float:
+        return self.service_snapshot_size(
+            get_service_by_id(snapshot.service_name), snapshot.id
+        )
diff --git a/selfprivacy_api/backup/backuper.py b/selfprivacy_api/backup/backuper.py
index f4c25a8..5d9b1c3 100644
--- a/selfprivacy_api/backup/backuper.py
+++ b/selfprivacy_api/backup/backuper.py
@@ -25,3 +25,7 @@ class AbstractBackuper(ABC):
     def restore_from_backup(self, repo_name: str, snapshot_id: str, folder: str):
         """Restore a target folder using a snapshot"""
         raise NotImplementedError
+
+    @abstractmethod
+    def restored_size(self, repo_name, snapshot_id) -> float:
+        raise NotImplementedError
diff --git a/selfprivacy_api/backup/restic_backuper.py b/selfprivacy_api/backup/restic_backuper.py
index 04461ca..a4a4830 100644
--- a/selfprivacy_api/backup/restic_backuper.py
+++ b/selfprivacy_api/backup/restic_backuper.py
@@ -90,6 +90,25 @@ class ResticBackuper(AbstractBackuper):
             if not "created restic repository" in output:
                 raise ValueError("cannot init a repo: " + output)
 
+    def restored_size(self, repo_name, snapshot_id) -> float:
+        """
+        Size of a snapshot
+        """
+        command = self.restic_command(
+            repo_name,
+            "stats",
+            snapshot_id,
+            "--json",
+        )
+
+        with subprocess.Popen(command, stdout=subprocess.PIPE, shell=False) as handle:
+            output = handle.communicate()[0].decode("utf-8")
+            try:
+                parsed_output = self.parse_json_output(output)
+                return parsed_output["total_size"]
+            except ValueError as e:
+                raise ValueError("cannot restore a snapshot: " + output) from e
+
     def restore_from_backup(self, repo_name, snapshot_id, folder):
         """
         Restore from backup with restic
@@ -135,7 +154,7 @@ class ResticBackuper(AbstractBackuper):
         if "Is there a repository at the following location?" in output:
             raise ValueError("No repository! : " + output)
         try:
-            return self.parse_snapshot_output(output)
+            return self.parse_json_output(output)
         except ValueError as e:
             raise ValueError("Cannot load snapshots: ") from e
 
@@ -152,10 +171,18 @@ class ResticBackuper(AbstractBackuper):
             snapshots.append(snapshot)
         return snapshots
 
-    def parse_snapshot_output(self, output: str) -> object:
-        if "[" not in output:
+    def parse_json_output(self, output: str) -> object:
+        indices = [
+            output.find("["),
+            output.find("{"),
+        ]
+        indices = [x for x in indices if x != -1]
+
+        if indices == []:
             raise ValueError(
                 "There is no json in the restic snapshot output : " + output
             )
-        starting_index = output.find("[")
+
+        starting_index = min(indices)
+
         return json.loads(output[starting_index:])
diff --git a/tests/test_graphql/test_backup.py b/tests/test_graphql/test_backup.py
index 0e8e246..4c6b2dd 100644
--- a/tests/test_graphql/test_backup.py
+++ b/tests/test_graphql/test_backup.py
@@ -127,3 +127,11 @@ def test_restore(backups, dummy_service):
 
     backups.restore_service_from_snapshot(dummy_service, snap.id)
     assert path.exists(path_to_nuke)
+
+
+def test_sizing(backups, dummy_service):
+    backups.back_up(dummy_service)
+    snap = backups.get_snapshots(dummy_service)[0]
+    size = backups.service_snapshot_size(dummy_service, snap.id)
+    assert size is not None
+    assert size > 0