From db407cd05081230fca5755c3a38c5944650f3bd2 Mon Sep 17 00:00:00 2001
From: Zichao Zhang <houbai.zzc@alibaba-inc.com>
Date: Thu, 18 Nov 2021 15:04:37 +0800
Subject: [PATCH] Support restoring from a specified backup (#561)

* Support restoring from specified backup

* Refactor backup test
---
 src/backup.rs        | 44 ++++++++++++++++++++++++++++++++++++++
 tests/test_backup.rs | 50 ++++++++++++++++++++++++++++++++++++++++----
 2 files changed, 90 insertions(+), 4 deletions(-)

diff --git a/src/backup.rs b/src/backup.rs
index d566eeb..75aed70 100644
--- a/src/backup.rs
+++ b/src/backup.rs
@@ -170,6 +170,50 @@ impl BackupEngine {
         Ok(())
     }
 
+    /// Restore from a specified backup
+    ///
+    /// The specified backup id should be passed in as an additional parameter.
+    pub fn restore_from_backup<D: AsRef<Path>, W: AsRef<Path>>(
+        &mut self,
+        db_dir: D,
+        wal_dir: W,
+        opts: &RestoreOptions,
+        backup_id: u32,
+    ) -> Result<(), Error> {
+        let db_dir = db_dir.as_ref();
+        let c_db_dir = if let Ok(c) = CString::new(db_dir.to_string_lossy().as_bytes()) {
+            c
+        } else {
+            return Err(Error::new(
+                "Failed to convert db_dir to CString \
+                     when restoring from latest backup"
+                    .to_owned(),
+            ));
+        };
+
+        let wal_dir = wal_dir.as_ref();
+        let c_wal_dir = if let Ok(c) = CString::new(wal_dir.to_string_lossy().as_bytes()) {
+            c
+        } else {
+            return Err(Error::new(
+                "Failed to convert wal_dir to CString \
+                     when restoring from latest backup"
+                    .to_owned(),
+            ));
+        };
+
+        unsafe {
+            ffi_try!(ffi::rocksdb_backup_engine_restore_db_from_backup(
+                self.inner,
+                c_db_dir.as_ptr(),
+                c_wal_dir.as_ptr(),
+                opts.inner,
+                backup_id,
+            ));
+        }
+        Ok(())
+    }
+
     /// Checks that each file exists and that the size of the file matches our
     /// expectations. it does not check file checksum.
     ///
diff --git a/tests/test_backup.rs b/tests/test_backup.rs
index c3e4885..3c945a4 100644
--- a/tests/test_backup.rs
+++ b/tests/test_backup.rs
@@ -23,17 +23,17 @@ use rocksdb::{
 use util::DBPath;
 
 #[test]
-fn backup_restore() {
+fn restore_from_latest() {
     // create backup
-    let path = DBPath::new("backup_test");
-    let restore_path = DBPath::new("restore_from_backup_path");
+    let path = DBPath::new("restore_from_latest_test");
+    let restore_path = DBPath::new("restore_from_latest_path");
     {
         let db = DB::open_default(&path).unwrap();
         assert!(db.put(b"k1", b"v1111").is_ok());
         let value = db.get(b"k1");
         assert_eq!(value.unwrap().unwrap(), b"v1111");
         {
-            let backup_path = DBPath::new("backup_path");
+            let backup_path = DBPath::new("backup_path_1");
             let backup_opts = BackupEngineOptions::default();
             let mut backup_engine = BackupEngine::open(&backup_opts, &backup_path).unwrap();
             assert!(backup_engine.create_new_backup(&db).is_ok());
@@ -61,3 +61,45 @@ fn backup_restore() {
         }
     }
 }
+
+#[test]
+fn restore_from_backup() {
+    // create backup
+    let path = DBPath::new("restore_from_backup_test");
+    let restore_path = DBPath::new("restore_from_backup_path");
+    {
+        let db = DB::open_default(&path).unwrap();
+        assert!(db.put(b"k1", b"v1111").is_ok());
+        let value = db.get(b"k1");
+        assert_eq!(value.unwrap().unwrap(), b"v1111");
+        {
+            let backup_path = DBPath::new("backup_path_2");
+            let backup_opts = BackupEngineOptions::default();
+            let mut backup_engine = BackupEngine::open(&backup_opts, &backup_path).unwrap();
+            assert!(backup_engine.create_new_backup(&db).is_ok());
+
+            // check backup info
+            let info = backup_engine.get_backup_info();
+            assert!(!info.is_empty());
+            info.iter().for_each(|i| {
+                assert!(backup_engine.verify_backup(i.backup_id).is_ok());
+                assert!(i.size > 0);
+            });
+
+            let backup_id = info.get(0).unwrap().backup_id;
+            let mut restore_option = RestoreOptions::default();
+            restore_option.set_keep_log_files(false); // true to keep log files
+            let restore_status = backup_engine.restore_from_backup(
+                &restore_path,
+                &restore_path,
+                &restore_option,
+                backup_id,
+            );
+            assert!(restore_status.is_ok());
+
+            let db_restore = DB::open_default(&restore_path).unwrap();
+            let value = db_restore.get(b"k1");
+            assert_eq!(value.unwrap().unwrap(), b"v1111");
+        }
+    }
+}