updated

Verticaleap · Verticaleap · commit 3209a63b1fb6 · 2021-06-07T15:49:28.000-05:00
diff --git a/src/db.rs b/src/db.rs
@@ -15,19 +15,21 @@ use tokio::sync::Semaphore;
 #[derive(Clone, Debug)]
 pub struct DB {
     pub client: Client,
-    pub db: String
+    pub db: String,
+    pub renamedb: Option<String>
 }
 
 type BoxResult<T> = std::result::Result<T, Box<dyn error::Error + Send + Sync>>;
 
 impl DB {
-    pub async fn init(url: &str, db: &str) -> BoxResult<Self> {
+    pub async fn init(url: &str, db: &str, renamedb: Option<&str>) -> BoxResult<Self> {
         let mut client_options = ClientOptions::parse(url).await?;
         client_options.app_name = Some("mongodb-stream-rs".to_string());
         client_options.read_concern = Some(ReadConcern::local());
         Ok(Self {
             client: Client::with_options(client_options)?,
-            db: db.to_owned()
+            db: db.to_owned(),
+            renamedb: renamedb.map(|s| s.into())
         })
     }
 
@@ -113,7 +115,7 @@ impl DB {
         };
 
         // Set counter to total
-        counter.total(total);
+        counter.set_total(total);
         
         let cursor = collection_handle.find(query, find_options).await?;
 
@@ -122,7 +124,10 @@ impl DB {
 
     pub async fn insert_cursor(&mut self, collection: &str, mut cursor: Cursor, mut counter: Counter) -> BoxResult<()> {
         // Get handle on collection
-        let coll = self.client.database(&self.db).collection(collection);
+        let coll = match &self.renamedb {
+            Some(db) => self.client.database(&db).collection(collection),
+            None => self.client.database(&self.db).collection(collection)
+        };
 
         log::info!("{}.{}: Inserting {} docs", self.db, collection, counter.total);
 
@@ -155,7 +160,10 @@ impl DB {
 
     pub async fn validate_docs(&mut self, collection: &str, mut cursor: Cursor, mut counter: Counter) -> BoxResult<()> {
         // Get handle on collection
-        let coll = self.client.database(&self.db).collection(collection);
+        let coll = match &self.renamedb {
+            Some(db) => self.client.database(&db).collection(collection),
+            None => self.client.database(&self.db).collection(collection)
+        };
 
         log::info!("{}.{}: Validating that {} docs in destination exist in source", self.db, collection, counter.total);
 
@@ -386,14 +394,18 @@ impl Counter {
 //        self.count = count as f64;
 //    }
 
-    pub fn total(&mut self, total: f64) {
+    pub fn set_total(&mut self, total: f64) {
         self.total = total;
     }
 
     pub fn count(&self) -> f64 {
         self.count
     }
 
+    pub fn total(&self) -> f64 {
+        self.total
+    }
+
     pub fn incr(&mut self, db: &str, collection: &str, count: f64, start: i64) {
         self.count += count;
         let percent = self.count / self.total * 100.0;
diff --git a/src/main.rs b/src/main.rs
@@ -117,6 +117,15 @@ async fn main() -> BoxResult<()> {
                 .help("Enable extra verbosity")
                 .takes_value(false)
         )
+        .arg(
+            Arg::with_name("rename")
+                .long("rename")
+                .required(false)
+                .value_name("MONGODB_RENAME")
+                .env("MONGODB_RENAME")
+                .help("Rename DB at destination")
+                .takes_value(true)
+        )
         .get_matches();
 
     // Initialize log Builder
@@ -139,15 +148,16 @@ async fn main() -> BoxResult<()> {
     let source = &opts.value_of("source_uri").unwrap();
     let destination= &opts.value_of("destination_uri").unwrap();
     let db = &opts.value_of("db").unwrap();
+    let renamedb = &opts.value_of("renamedb");
 
     println!(
         "Starting mongodb-stream-rs:{}", 
         crate_version!(),
     );
 
     // Create connections to source and destination db's
-    let source_db = DB::init(&source, &db).await?;
-    let destination_db = DB::init(&destination, &db).await?;
+    let source_db = DB::init(&source, &db, None).await?;
+    let destination_db = DB::init(&destination, &db, *renamedb).await?;
 
 
     let collections = match &opts.is_present("collection") {
