From f192c39be8e103b42656422e75c4a5518cab126c Mon Sep 17 00:00:00 2001
From: simon <simobilleter@gmail.com>
Date: Sun, 27 Nov 2022 18:01:57 +0700
Subject: [PATCH] make snapshot before mapping changes

---
 tubearchivist/home/src/es/index_setup.py | 11 ++++++++-
 tubearchivist/home/src/es/snapshot.py    | 30 +++++++++++++++++++++++-
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/tubearchivist/home/src/es/index_setup.py b/tubearchivist/home/src/es/index_setup.py
index 64cc9ed..c913da7 100644
--- a/tubearchivist/home/src/es/index_setup.py
+++ b/tubearchivist/home/src/es/index_setup.py
@@ -7,6 +7,8 @@ functionality:
 
 from home.src.es.backup import ElasticBackup
 from home.src.es.connect import ElasticWrap
+from home.src.es.snapshot import ElasticSnapshot
+from home.src.ta.config import AppConfig
 from home.src.ta.helper import get_mapping
 
 
@@ -197,5 +199,12 @@ class ElasitIndexWrap:
         if self.backup_run:
             return
 
-        ElasticBackup(reason="update").backup_all_indexes()
+        config = AppConfig().config
+        if config["application"]["enable_snapshot"]:
+            # take snapshot if enabled
+            ElasticSnapshot().take_snapshot_now(wait=True)
+        else:
+            # fallback to json backup
+            ElasticBackup(reason="update").backup_all_indexes()
+
         self.backup_run = True
diff --git a/tubearchivist/home/src/es/snapshot.py b/tubearchivist/home/src/es/snapshot.py
index 19e49e2..48baeaf 100644
--- a/tubearchivist/home/src/es/snapshot.py
+++ b/tubearchivist/home/src/es/snapshot.py
@@ -5,6 +5,7 @@ functionality:
 
 from datetime import datetime
 from os import environ
+from time import sleep
 from zoneinfo import ZoneInfo
 
 from home.src.es.connect import ElasticWrap
@@ -142,15 +143,42 @@ class ElasticSnapshot:
         print("snapshot: last snapshot is up-to-date")
         return outdated
 
-    def take_snapshot_now(self):
+    def take_snapshot_now(self, wait=False):
         """execute daily snapshot now"""
         path = f"_slm/policy/{self.POLICY}/_execute"
         response, statuscode = ElasticWrap(path).post()
         if statuscode == 200:
             print(f"snapshot: executing now: {response}")
 
+        if wait:
+            self._wait_for_snapshot(response["snapshot_name"])
+
         return response
 
+    def _wait_for_snapshot(self, snapshot_name):
+        """return after snapshot_name completes"""
+        path = f"_snapshot/{self.REPO}/{snapshot_name}"
+
+        while True:
+            # wait for task to be created
+            sleep(1)
+            _, statuscode = ElasticWrap(path).get()
+            if statuscode == 200:
+                break
+
+        while True:
+            # wait for snapshot success
+            response, statuscode = ElasticWrap(path).get()
+            snapshot_state = response["snapshots"][0]["state"]
+            if snapshot_state == "SUCCESS":
+                break
+
+            print(f"snapshot: {snapshot_name} in state {snapshot_state}")
+            print("snapshot: wait to complete")
+            sleep(5)
+
+        print(f"snapshot: completed - {response}")
+
     def get_snapshot_stats(self):
         """get snapshot info for frontend"""
         snapshot_info = self._build_policy_details()