From eb53ed4cc14312e65e389c8a0183547b1932e15c Mon Sep 17 00:00:00 2001
From: mpostma <postma.marin@protonmail.com>
Date: Tue, 23 Mar 2021 16:37:46 +0100
Subject: [PATCH] load snapshot

---
 meilisearch-http/src/helpers/compression.rs   |  6 ++--
 meilisearch-http/src/index_controller/mod.rs  | 10 ++++--
 .../src/index_controller/snapshot.rs          | 33 +++++++++++++++++--
 3 files changed, 41 insertions(+), 8 deletions(-)

diff --git a/meilisearch-http/src/helpers/compression.rs b/meilisearch-http/src/helpers/compression.rs
index 74578ed0c..201f35149 100644
--- a/meilisearch-http/src/helpers/compression.rs
+++ b/meilisearch-http/src/helpers/compression.rs
@@ -5,9 +5,7 @@ use std::path::Path;
 use flate2::{Compression, write::GzEncoder, read::GzDecoder};
 use tar::{Archive, Builder};
 
-use crate::error::Error;
-
-pub fn to_tar_gz(src: impl AsRef<Path>, dest: impl AsRef<Path>) -> Result<(), Error> {
+pub fn to_tar_gz(src: impl AsRef<Path>, dest: impl AsRef<Path>) -> anyhow::Result<()> {
     let mut f = File::create(dest)?;
     let gz_encoder = GzEncoder::new(&mut f, Compression::default());
     let mut tar_encoder = Builder::new(gz_encoder);
@@ -18,7 +16,7 @@ pub fn to_tar_gz(src: impl AsRef<Path>, dest: impl AsRef<Path>) -> Result<(), Er
     Ok(())
 }
 
-pub fn from_tar_gz(src: impl AsRef<Path>, dest: impl AsRef<Path>) -> Result<(), Error> {
+pub fn from_tar_gz(src: impl AsRef<Path>, dest: impl AsRef<Path>) -> anyhow::Result<()> {
     let f = File::open(&src)?;
     let gz = GzDecoder::new(f);
     let mut ar = Archive::new(gz);
diff --git a/meilisearch-http/src/index_controller/mod.rs b/meilisearch-http/src/index_controller/mod.rs
index bcd0cd8fe..73d071daa 100644
--- a/meilisearch-http/src/index_controller/mod.rs
+++ b/meilisearch-http/src/index_controller/mod.rs
@@ -20,10 +20,11 @@ use tokio::time::sleep;
 use crate::index::{Document, SearchQuery, SearchResult};
 use crate::index::{Facets, Settings, UpdateResult};
 use crate::option::Opt;
-use crate::helpers::compression;
+
 use index_actor::IndexActorHandle;
 use update_actor::UpdateActorHandle;
 use uuid_resolver::UuidResolverHandle;
+use snapshot::load_snapshot;
 
 use snapshot::SnapshotService;
 pub use updates::{Failed, Processed, Processing};
@@ -71,7 +72,12 @@ impl IndexController {
         let update_store_size = options.max_udb_size.get_bytes() as usize;
 
         if let Some(ref path) = options.import_snapshot {
-            compression::from_tar_gz(path, &options.db_path)?;
+            load_snapshot(
+                &options.db_path,
+                path,
+                options.ignore_snapshot_if_db_exists,
+                options.ignore_missing_snapshot,
+            )?;
         }
 
         let uuid_resolver = uuid_resolver::UuidResolverHandleImpl::new(&path)?;
diff --git a/meilisearch-http/src/index_controller/snapshot.rs b/meilisearch-http/src/index_controller/snapshot.rs
index 998821d9f..586d096ba 100644
--- a/meilisearch-http/src/index_controller/snapshot.rs
+++ b/meilisearch-http/src/index_controller/snapshot.rs
@@ -1,4 +1,4 @@
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
 use std::time::Duration;
 
 use anyhow::bail;
@@ -93,9 +93,38 @@ where
     }
 }
 
+pub fn load_snapshot(
+    db_path: impl AsRef<Path>,
+    snapshot_path: impl AsRef<Path>,
+    ignore_snapshot_if_db_exists: bool,
+    ignore_missing_snapshot: bool,
+) -> anyhow::Result<()> {
+    if !db_path.as_ref().exists() && snapshot_path.as_ref().exists() {
+        compression::from_tar_gz(snapshot_path, db_path)
+    } else if db_path.as_ref().exists() && !ignore_snapshot_if_db_exists {
+        bail!(
+            "database already exists at {:?}, try to delete it or rename it",
+            db_path
+                .as_ref()
+                .canonicalize()
+                .unwrap_or(db_path.as_ref().to_owned())
+        )
+    } else if !snapshot_path.as_ref().exists() && !ignore_missing_snapshot {
+        bail!(
+            "snapshot doesn't exist at {:?}",
+            snapshot_path
+                .as_ref()
+                .canonicalize()
+                .unwrap_or(snapshot_path.as_ref().to_owned())
+        )
+    } else {
+        Ok(())
+    }
+}
+
 #[cfg(test)]
 mod test {
-    use futures::future::{ok, err};
+    use futures::future::{err, ok};
     use rand::Rng;
     use tokio::time::timeout;
     use uuid::Uuid;