database backup and restore

2 files changed, 49 insertions, 5 deletions

diff --git a/src/bin/client.rs b/src/bin/client.rs
index a91c2dc..9f5fbfc 100644
--- a/src/bin/client.rs
+++ b/src/bin/client.rs
@@ -34,7 +34,7 @@ fn main() {
     let server;
     let mut name = String::new();
 
-    let matches = App::new("space")
+    let matches = App::new("space client")
         .subcommand(SubCommand::with_name("mining"))
         .subcommand(SubCommand::with_name("engines"))
         .subcommand(SubCommand::with_name("refinery"))
diff --git a/src/bin/server.rs b/src/bin/server.rs
index 881cc4d..708cd3e 100644
--- a/src/bin/server.rs
+++ b/src/bin/server.rs
@@ -1,14 +1,19 @@
 extern crate space;
 
+use clap::{App, SubCommand};
+use diesel::pg::PgConnection;
+use diesel::prelude::*;
 use std::collections::HashMap;
 use std::io::Write;
 use std::net::TcpListener;
-use std::thread::sleep;
+use std::thread::{sleep, spawn};
 use std::time::{Duration, Instant};
 
 use space::constants;
-use space::mass::Mass;
-use space::math::rand_name;
+use space::mass::{Mass, MassEntry};
+use space::math::{get_db_url, rand_name};
+use space::schema::masses::dsl::masses as db_masses;
+use space::schema::masses::dsl::name as name_column;
 use space::server_connection::ServerConnection;
 
 fn populate() -> HashMap<String, Mass> {
@@ -21,12 +26,44 @@ fn populate() -> HashMap<String, Mass> {
     masses
 }
 
+fn backup(masses: HashMap<String, Mass>) {
+    let connection = PgConnection::establish(&get_db_url()).expect("Cannot connect");
+    for (name, mass) in masses {
+        let mass_entry = mass.to_mass_entry(name.to_string());
+        diesel::insert_into(db_masses)
+            .values(&mass_entry)
+            .on_conflict(name_column)
+            .do_update()
+            .set(&mass_entry)
+            .execute(&connection)
+            .expect("Cannot backup");
+    }
+}
+
+fn restore() -> HashMap<String, Mass> {
+    let connection = PgConnection::establish(&get_db_url()).expect("Cannot connect");
+    db_masses
+        .load::<MassEntry>(&connection)
+        .expect("Cannot query, are you sure you can restore?")
+        .iter()
+        .map(|mass_entry| mass_entry.to_mass())
+        .collect()
+}
+
 fn main() {
     let listener = TcpListener::bind("localhost:6000").unwrap();
     listener.set_nonblocking(true).unwrap();
 
-    let mut masses = populate();
+    let matches = App::new("space server")
+        .subcommand(SubCommand::with_name("--restore"))
+        .get_matches();
+
+    let mut masses = match matches.subcommand_name() {
+        Some("--restore") => restore(),
+        _ => populate(),
+    };
 
+    let mut backup_countdown = constants::BACKUP_COUNTDOWN;
     let mut connections: Vec<ServerConnection> = Vec::new();
     for stream in listener.incoming() {
         match stream {
@@ -64,11 +101,18 @@ fn main() {
                     masses.insert(key.to_string(), mass);
                 }
 
+                if backup_countdown == 0 {
+                    let masses_clone = masses.clone();
+                    spawn(move || backup(masses_clone));
+                    backup_countdown = constants::BACKUP_COUNTDOWN;
+                }
+
                 if timer.elapsed().as_millis() < constants::LOOP_DURATION_MS.into() {
                     sleep(Duration::from_millis(
                         constants::LOOP_DURATION_MS - timer.elapsed().as_millis() as u64,
                     ));
                 }
+                backup_countdown -= 1;
             }
         }
     }