From d7899e0cb4a6683af6c9a9b6755894c76f177f55 Mon Sep 17 00:00:00 2001
From: Ziyang Hu <hu.ziyang@cantab.net>
Date: Thu, 24 Nov 2022 17:15:03 +0800
Subject: [PATCH] cozoserver restore logic

---
 cozo-core/src/runtime/db.rs |  2 +-
 cozoserver/src/main.rs      | 36 ++++++++++--------------------------
 2 files changed, 11 insertions(+), 27 deletions(-)

diff --git a/cozo-core/src/runtime/db.rs b/cozo-core/src/runtime/db.rs
index b4fded0a..f9836879 100644
--- a/cozo-core/src/runtime/db.rs
+++ b/cozo-core/src/runtime/db.rs
@@ -372,7 +372,7 @@ impl<'s, S: Storage<'s>> Db<S> {
         {
             let sqlite_db = crate::new_cozo_sqlite(in_file)?;
             let s_tx = sqlite_db.transact_write()?;
-            let store_id = s_tx.relation_store_id.load(Ordering::AcqRel);
+            let store_id = s_tx.relation_store_id.load(Ordering::SeqCst);
             if store_id != 0 {
                 bail!(
                     "Cannot restore backup: data exists in the current database. \
diff --git a/cozoserver/src/main.rs b/cozoserver/src/main.rs
index b11b3098..629b8569 100644
--- a/cozoserver/src/main.rs
+++ b/cozoserver/src/main.rs
@@ -24,16 +24,19 @@ use cozo::*;
 #[clap(version, about, long_about = None)]
 struct Args {
     /// Database kind, can be `mem`, `sqlite`, `rocksdb` and others.
-    /// Some kinds may not be available if the executable did not set its build option.
     #[clap(short, long, default_value_t = String::from("mem"))]
     kind: String,
 
     /// Path to the directory to store the database
-    #[clap(short, long, default_value_t = String::from(""))]
+    #[clap(short, long, default_value_t = String::from("cozo.db"))]
     path: String,
 
+    /// Restore from the specified backup before starting the server
+    #[clap(long)]
+    restore: Option<String>,
+
     /// Extra config in JSON format
-    #[clap(short, long, default_value_t = serde_json::Value::Null)]
+    #[clap(short, long, default_value_t = json!({}))]
     config: serde_json::Value,
 
     /// Address to bind the service to
@@ -67,6 +70,10 @@ fn main() {
 
     let db = DbInstance::new(args.kind.as_str(), args.path.as_str(), args.config.clone()).unwrap();
 
+    if let Some(restore_path) = &args.restore {
+        db.restore_backup(restore_path.to_string()).unwrap();
+    }
+
     let conf_path = format!("{}.{}.cozo_auth", args.path, args.kind);
     let auth_guard = match fs::read_to_string(&conf_path) {
         Ok(s) => s.trim().to_string(),
@@ -205,29 +212,6 @@ fn main() {
                         }
                     }
                 },
-                (POST) (/restore) => {
-                    check_auth!(request, auth_guard);
-
-                    #[derive(serde_derive::Serialize, serde_derive::Deserialize)]
-                    struct BackupPayload {
-                        path: String,
-                    }
-
-                    let payload: BackupPayload = try_or_400!(rouille::input::json_input(request));
-
-                    let result = db.restore_backup(payload.path.clone());
-
-                    match result {
-                        Ok(()) => {
-                            let ret = json!({"ok": true});
-                            Response::json(&ret)
-                        }
-                        Err(err) => {
-                            let ret = json!({"ok": false, "message": err.to_string()});
-                            Response::json(&ret).with_status_code(400)
-                        }
-                    }
-                },
                 (GET) (/) => {
                     Response::html(HTML_CONTENT)
                 },