Add Buffer Size Parameter to dump create

replibyte/src/cli.rs

@@ -104,6 +104,8 @@ pub struct DumpCreateArgs {
     /// import dump from stdin
     #[clap(name = "input", short, long, requires = "source_type")]
     pub input: bool,
+    #[clap(name = "dump_chunk_size", short, long, default_value_t = 100, value_name = "dump_chunk_size", help = "Dump Chunk Size in MB")]
+    pub dump_chunk_size: usize,
     #[clap(short, long, parse(from_os_str), value_name = "dump file")]
     /// dump file
     pub file: Option<PathBuf>,

replibyte/src/commands/dump.rs

@@ -132,6 +132,7 @@ where
                 skip_config: &skip_config,
                 database_subset: &source.database_subset,
                 only_tables: &only_tables_config,
+                dump_chunk_size: &args.dump_chunk_size,
             };
 
             match args.source_type.as_ref().map(|x| x.as_str()) {

replibyte/src/source/mod.rs

@@ -29,4 +29,5 @@ pub struct SourceOptions<'a> {
     pub skip_config: &'a Vec<SkipConfig>,
     pub database_subset: &'a Option<DatabaseSubsetConfig>,
     pub only_tables: &'a Vec<OnlyTablesConfig>,
+    pub dump_chunk_size: &'a usize,
 }

replibyte/src/source/mongodb.rs

@@ -357,6 +357,10 @@ mod tests {
         )
     }
 
+    fn get_default_dump_chunk_size() -> usize {
+        100
+    }
+
     #[test]
     fn connect() {
         let p = get_mongodb();
@@ -368,6 +372,7 @@ mod tests {
             skip_config: &vec![],
             database_subset: &None,
             only_tables: &vec![],
+            dump_chunk_size: &get_default_dump_chunk_size(),
         };
 
         assert!(p.read(source_options, |_, _| {}).is_ok());
@@ -380,6 +385,7 @@ mod tests {
             skip_config: &vec![],
             database_subset: &None,
             only_tables: &vec![],
+            dump_chunk_size: &get_default_dump_chunk_size(),
         };
 
         assert!(p.read(source_options, |_, _| {}).is_err());
@@ -395,6 +401,7 @@ mod tests {
             skip_config: &vec![],
             database_subset: &None,
             only_tables: &vec![],
+            dump_chunk_size: &get_default_dump_chunk_size(),
         };
 
         p.read(source_options, |original_query, query| {

replibyte/src/source/mysql.rs

@@ -444,6 +444,10 @@ mod tests {
         Mysql::new("127.0.0.1", 3306, "world", "root", "wrong_password")
     }
 
+    fn get_default_dump_chunk_size() -> usize {
+        100
+    }
+
     #[test]
     fn connect() {
         let mut p = get_mysql();
@@ -456,6 +460,7 @@ mod tests {
             skip_config: &vec![],
             database_subset: &None,
             only_tables: &vec![],
+            dump_chunk_size: &get_default_dump_chunk_size(),
         };
 
         assert!(p.read(source_options, |_original_query, _query| {}).is_ok());
@@ -468,6 +473,7 @@ mod tests {
             skip_config: &vec![],
             database_subset: &None,
             only_tables: &vec![],
+            dump_chunk_size: &get_default_dump_chunk_size(),
         };
         assert!(p
             .read(source_options, |_original_query, _query| {})
@@ -484,6 +490,7 @@ mod tests {
             skip_config: &vec![],
             database_subset: &None,
             only_tables: &vec![],
+            dump_chunk_size: &get_default_dump_chunk_size(),
         };
         let _ = p.read(source_options, |original_query, query| {
             assert!(original_query.data().len() > 0);

replibyte/src/source/postgres.rs

@@ -571,6 +571,10 @@ mod tests {
         Postgres::new("localhost", 5432, "root", "root", "wrongpassword")
     }
 
+    fn get_default_dump_chunk_size() -> usize {
+        100
+    }
+
     #[test]
     fn connect() {
         let p = get_postgres();
@@ -581,6 +585,7 @@ mod tests {
             skip_config: &vec![],
             database_subset: &None,
             only_tables: &vec![],
+            dump_chunk_size: &get_default_dump_chunk_size(),
         };
 
         assert!(p.read(source_options, |original_query, query| {}).is_ok());
@@ -593,6 +598,7 @@ mod tests {
             skip_config: &vec![],
             database_subset: &None,
             only_tables: &vec![],
+            dump_chunk_size: &get_default_dump_chunk_size(),
         };
 
         assert!(p.read(source_options, |original_query, query| {}).is_err());
@@ -608,6 +614,7 @@ mod tests {
             skip_config: &vec![],
             database_subset: &None,
             only_tables: &vec![],
+            dump_chunk_size: &get_default_dump_chunk_size(),
         };
 
         let _ = p.read(source_options, |original_query, query| {
@@ -734,6 +741,7 @@ mod tests {
             skip_config: &vec![],
             database_subset: &None,
             only_tables: &vec![],
+            dump_chunk_size: &get_default_dump_chunk_size(),
         };
 
         let _ = p.read(source_options, |original_query, query| {
@@ -775,6 +783,7 @@ mod tests {
             skip_config: &skip_config,
             database_subset: &None,
             only_tables: &vec![],
+            dump_chunk_size: &get_default_dump_chunk_size(),
         };
 
         let _ = p.read(source_options, |_original_query, query| {
@@ -826,6 +835,7 @@ mod tests {
                 passthrough_tables: None,
             }),
             only_tables: &vec![],
+            dump_chunk_size: &get_default_dump_chunk_size(),
         };
 
         let mut rows_percent_50 = vec![];
@@ -863,6 +873,7 @@ mod tests {
                 passthrough_tables: None,
             }),
             only_tables: &vec![],
+            dump_chunk_size: &get_default_dump_chunk_size(),
         };
 
         let mut rows_percent_30 = vec![];

replibyte/src/tasks/full_dump.rs

@@ -70,7 +70,7 @@ where
         });
 
         // buffer of 100MB in memory to use and re-use to upload data into datastore
-        let buffer_size = 100 * 1024 * 1024;
+        let buffer_size = self.options.dump_chunk_size * 1024 * 1024;
         let mut queries = vec![];
         let mut consumed_buffer_size = 0usize;
         let mut total_transferred_bytes = 0usize;