From 6da0791b4787f62dd1e104f5ed0f686a9033d392 Mon Sep 17 00:00:00 2001
From: Houkime <>
Date: Wed, 7 Jun 2023 15:05:58 +0000
Subject: [PATCH] feature(backups): integration between restore and jobs

---
 selfprivacy_api/backup/__init__.py | 26 ++++++++++++++++++++++----
 selfprivacy_api/backup/jobs.py     | 19 +++++++++++++++++++
 selfprivacy_api/backup/tasks.py    |  5 ++++-
 3 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/selfprivacy_api/backup/__init__.py b/selfprivacy_api/backup/__init__.py
index 7678258..b313165 100644
--- a/selfprivacy_api/backup/__init__.py
+++ b/selfprivacy_api/backup/__init__.py
@@ -13,7 +13,12 @@ from selfprivacy_api.graphql.queries.providers import BackupProvider
 from selfprivacy_api.backup.providers.provider import AbstractBackupProvider
 from selfprivacy_api.backup.providers import get_provider
 from selfprivacy_api.backup.storage import Storage
-from selfprivacy_api.backup.jobs import get_backup_job, add_backup_job
+from selfprivacy_api.backup.jobs import (
+    get_backup_job,
+    add_backup_job,
+    get_restore_job,
+    add_restore_job,
+)
 from selfprivacy_api.jobs import Jobs, JobStatus
 
 
@@ -285,6 +290,7 @@ class Backups:
         for snapshot in upstream_snapshots:
             Storage.cache_snapshot(snapshot)
 
+    # to be deprecated/internalized in favor of restore_snapshot()
     @staticmethod
     def restore_service_from_snapshot(service: Service, snapshot_id: str):
         repo_name = service.get_id()
@@ -294,9 +300,21 @@ class Backups:
 
     @staticmethod
     def restore_snapshot(snapshot: Snapshot):
-        Backups.restore_service_from_snapshot(
-            get_service_by_id(snapshot.service_name), snapshot.id
-        )
+        service = get_service_by_id(snapshot.service_name)
+
+        job = get_restore_job(service)
+        if job is None:
+            job = add_restore_job(snapshot)
+
+        Jobs.update(job, status=JobStatus.RUNNING)
+        try:
+            Backups.restore_service_from_snapshot(service, snapshot.id)
+            service.post_restore()
+        except Exception as e:
+            Jobs.update(job, status=JobStatus.ERROR)
+            raise e
+
+        Jobs.update(job, status=JobStatus.FINISHED)
 
     @staticmethod
     def service_snapshot_size(service: Service, snapshot_id: str) -> float:
diff --git a/selfprivacy_api/backup/jobs.py b/selfprivacy_api/backup/jobs.py
index 2293da0..5a9cb0d 100644
--- a/selfprivacy_api/backup/jobs.py
+++ b/selfprivacy_api/backup/jobs.py
@@ -1,7 +1,9 @@
 from typing import Optional, List
 
+from selfprivacy_api.models.backup.snapshot import Snapshot
 from selfprivacy_api.jobs import Jobs, Job, JobStatus
 from selfprivacy_api.services.service import Service
+from selfprivacy_api.services import get_service_by_id
 
 
 def job_type_prefix(service: Service) -> str:
@@ -47,6 +49,23 @@ def add_backup_job(service: Service) -> Job:
     return job
 
 
+def add_restore_job(snapshot: Snapshot) -> Job:
+    service = get_service_by_id(snapshot.service_name)
+    if is_something_queued_for(service):
+        message = (
+            f"Cannot start a restore of {service.get_id()}, another operation is queued: "
+            + get_jobs_by_service(service)[0].type_id
+        )
+        raise ValueError(message)
+    display_name = service.get_display_name()
+    job = Jobs.add(
+        type_id=restore_job_type(service),
+        name=f"Restore {display_name}",
+        description=f"restoring {display_name} from {snapshot.id}",
+    )
+    return job
+
+
 def get_job_by_type(type_id: str) -> Optional[Job]:
     for job in Jobs.get_jobs():
         if job.type_id == type_id and job.status in [
diff --git a/selfprivacy_api/backup/tasks.py b/selfprivacy_api/backup/tasks.py
index d92a926..e88f651 100644
--- a/selfprivacy_api/backup/tasks.py
+++ b/selfprivacy_api/backup/tasks.py
@@ -2,9 +2,10 @@ from datetime import datetime
 
 from selfprivacy_api.models.backup.snapshot import Snapshot
 from selfprivacy_api.utils.huey import huey
+from selfprivacy_api.services import get_service_by_id
 from selfprivacy_api.services.service import Service
 from selfprivacy_api.backup import Backups
-from selfprivacy_api.backup.jobs import add_backup_job
+from selfprivacy_api.backup.jobs import add_backup_job, add_restore_job
 
 
 def validate_datetime(dt: datetime):
@@ -32,6 +33,8 @@ def start_backup(service: Service) -> bool:
 
 @huey.task()
 def restore_snapshot(snapshot: Snapshot) -> bool:
+    add_restore_job(snapshot)
+
     Backups.restore_snapshot(snapshot)
     return True