feat: added subcommands to aw-sync, support selectively syncing buckets

Author: ErikBjare

aw-sync/src/accessmethod.rs

@@ -0,0 +1,108 @@
+use std::collections::HashMap;
+
+use aw_client_rust::AwClient;
+use chrono::{DateTime, Utc};
+use reqwest::StatusCode;
+
+use aw_datastore::{Datastore, DatastoreError};
+use aw_models::{Bucket, Event};
+
+// This trait should be implemented by both AwClient and Datastore, unifying them under a single API
+pub trait AccessMethod: std::fmt::Debug {
+    fn get_buckets(&self) -> Result<HashMap<String, Bucket>, String>;
+    fn get_bucket(&self, bucket_id: &str) -> Result<Bucket, DatastoreError>;
+    fn create_bucket(&self, bucket: &Bucket) -> Result<(), DatastoreError>;
+    fn get_events(
+        &self,
+        bucket_id: &str,
+        start: Option<DateTime<Utc>>,
+        end: Option<DateTime<Utc>>,
+        limit: Option<u64>,
+    ) -> Result<Vec<Event>, String>;
+    fn insert_events(&self, bucket_id: &str, events: Vec<Event>) -> Result<Vec<Event>, String>;
+    fn get_event_count(&self, bucket_id: &str) -> Result<i64, String>;
+    fn heartbeat(&self, bucket_id: &str, event: Event, duration: f64) -> Result<(), String>;
+}
+
+impl AccessMethod for Datastore {
+    fn get_buckets(&self) -> Result<HashMap<String, Bucket>, String> {
+        Ok(self.get_buckets().unwrap())
+    }
+    fn get_bucket(&self, bucket_id: &str) -> Result<Bucket, DatastoreError> {
+        self.get_bucket(bucket_id)
+    }
+    fn create_bucket(&self, bucket: &Bucket) -> Result<(), DatastoreError> {
+        self.create_bucket(bucket)?;
+        self.force_commit().unwrap();
+        Ok(())
+    }
+    fn get_events(
+        &self,
+        bucket_id: &str,
+        start: Option<DateTime<Utc>>,
+        end: Option<DateTime<Utc>>,
+        limit: Option<u64>,
+    ) -> Result<Vec<Event>, String> {
+        Ok(self.get_events(bucket_id, start, end, limit).unwrap())
+    }
+    fn heartbeat(&self, bucket_id: &str, event: Event, duration: f64) -> Result<(), String> {
+        self.heartbeat(bucket_id, event, duration).unwrap();
+        self.force_commit().unwrap();
+        Ok(())
+    }
+    fn insert_events(&self, bucket_id: &str, events: Vec<Event>) -> Result<Vec<Event>, String> {
+        let res = self.insert_events(bucket_id, &events[..]).unwrap();
+        self.force_commit().unwrap();
+        Ok(res)
+    }
+    fn get_event_count(&self, bucket_id: &str) -> Result<i64, String> {
+        Ok(self.get_event_count(bucket_id, None, None).unwrap())
+    }
+}
+
+impl AccessMethod for AwClient {
+    fn get_buckets(&self) -> Result<HashMap<String, Bucket>, String> {
+        Ok(self.get_buckets().unwrap())
+    }
+    fn get_bucket(&self, bucket_id: &str) -> Result<Bucket, DatastoreError> {
+        let bucket = self.get_bucket(bucket_id);
+        match bucket {
+            Ok(bucket) => Ok(bucket),
+            Err(e) => {
+                warn!("{:?}", e);
+                let code = e.status().unwrap();
+                if code == StatusCode::NOT_FOUND {
+                    Err(DatastoreError::NoSuchBucket(bucket_id.into()))
+                } else {
+                    panic!("Unexpected error");
+                }
+            }
+        }
+    }
+    fn get_events(
+        &self,
+        bucket_id: &str,
+        start: Option<DateTime<Utc>>,
+        end: Option<DateTime<Utc>>,
+        limit: Option<u64>,
+    ) -> Result<Vec<Event>, String> {
+        Ok(self.get_events(bucket_id, start, end, limit).unwrap())
+    }
+    fn insert_events(&self, _bucket_id: &str, _events: Vec<Event>) -> Result<Vec<Event>, String> {
+        //Ok(self.insert_events(bucket_id, &events[..]).unwrap())
+        Err("Not implemented".to_string())
+    }
+    fn get_event_count(&self, bucket_id: &str) -> Result<i64, String> {
+        Ok(self.get_event_count(bucket_id).unwrap())
+    }
+    fn create_bucket(&self, bucket: &Bucket) -> Result<(), DatastoreError> {
+        self.create_bucket(bucket.id.as_str(), bucket._type.as_str())
+            .unwrap();
+        Ok(())
+        //Err(DatastoreError::InternalError("Not implemented".to_string()))
+    }
+    fn heartbeat(&self, bucket_id: &str, event: Event, duration: f64) -> Result<(), String> {
+        self.heartbeat(bucket_id, &event, duration)
+            .map_err(|e| format!("{:?}", e))
+    }
+}

aw-sync/src/lib.rs

@@ -7,3 +7,6 @@ extern crate serde_json;
 mod sync;
 pub use sync::sync_datastores;
 pub use sync::sync_run;
+
+mod accessmethod;
+pub use accessmethod::AccessMethod;

aw-sync/src/main.rs

@@ -16,17 +16,21 @@ extern crate serde_json;
 use std::path::Path;
 
 use chrono::{DateTime, Utc};
-use clap::Parser;
+use clap::{Parser, Subcommand};
 
 use aw_client_rust::AwClient;
 
+mod accessmethod;
 mod sync;
 
 const DEFAULT_PORT: &str = "5600";
 
 #[derive(Parser)]
 #[clap(version = "0.1", author = "Erik Bjäreholt")]
 struct Opts {
+    #[clap(subcommand)]
+    command: Commands,
+
     /// Host of instance to connect to
     #[clap(long, default_value = "127.0.0.1")]
     host: String,
@@ -40,16 +44,26 @@ struct Opts {
     /// If not specified, exit
     #[clap(long)]
     sync_dir: String,
-    /// Date to start syncing from
-    /// If not specified, start from beginning
-    /// NOTE: might be unstable, as count cannot be used to verify integrity of sync.
-    /// Format: YYYY-MM-DD
-    #[clap(long)]
-    start_date: Option<String>,
-    /// Specify buckets to sync
-    /// If not specified, all buckets will be synced
-    #[clap(long)]
-    buckets: Vec<String>,
+}
+
+#[derive(Subcommand)]
+enum Commands {
+    /// Clones repos
+    #[clap(arg_required_else_help = true)]
+    Sync {
+        /// Date to start syncing from
+        /// If not specified, start from beginning
+        /// NOTE: might be unstable, as count cannot be used to verify integrity of sync.
+        /// Format: YYYY-MM-DD
+        #[clap(long)]
+        start_date: Option<String>,
+        /// Specify buckets to sync
+        /// If not specified, all buckets will be synced
+        #[clap(long)]
+        buckets: Vec<String>,
+    },
+    /// List buckets and their sync status
+    List {},
 }
 
 fn main() -> std::io::Result<()> {
@@ -74,16 +88,37 @@ fn main() -> std::io::Result<()> {
 
     let client = AwClient::new(opts.host.as_str(), port, "aw-sync-rust");
 
-    let start: Option<DateTime<Utc>> = opts.start_date.map(|date| {
-        chrono::DateTime::parse_from_rfc3339(&date)
-            .unwrap()
-            .with_timezone(&chrono::Utc)
-    });
-    sync::sync_run(sync_directory, client, opts.buckets, start).map_err(|e| {
+    match &opts.command {
+        // Perform two-way sync
+        Commands::Sync {
+            start_date,
+            buckets,
+        } => {
+            let start: Option<DateTime<Utc>> = start_date.as_ref().map(|date| {
+                let date_copy = date.clone();
+                chrono::DateTime::parse_from_rfc3339(&date_copy)
+                    .unwrap()
+                    .with_timezone(&chrono::Utc)
+            });
+
+            sync::sync_run(sync_directory, client, buckets, start).map_err(|e| {
+                println!("Error: {}", e);
+                std::io::Error::new(std::io::ErrorKind::Other, e)
+            })?;
+            info!("Finished successfully, exiting...");
+            Ok(())
+        }
+
+        // List all buckets
+        Commands::List {} => {
+            sync::list_buckets(&client, sync_directory);
+            Ok(())
+        }
+    }
+    .map_err(|e: String| {
         println!("Error: {}", e);
         std::io::Error::new(std::io::ErrorKind::Other, e)
     })?;
-    info!("Finished successfully, exiting...");
 
     // Needed to give the datastores some time to commit before program is shut down.
     // 100ms isn't actually needed, seemed to work fine with as little as 10ms, but I'd rather give