Add option to specify chunk size

replibyte/src/cli.rs

@@ -110,6 +110,9 @@ pub struct DumpCreateArgs {
     /// dump name
     #[clap(short, long)]
     pub name: Option<String>,
+    /// chunk size, default to 100
+    #[clap(short, long)]
+    pub chunk_size: Option<usize>,
 }
 
 #[derive(Args, Debug)]

replibyte/src/commands/dump.rs

@@ -132,6 +132,7 @@ where
                 skip_config: &skip_config,
                 database_subset: &source.database_subset,
                 only_tables: &only_tables_config,
+                chunk_size: &args.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 chunk_size: &'a Option<usize>,
 }

replibyte/src/source/mongodb.rs

@@ -368,6 +368,7 @@ mod tests {
             skip_config: &vec![],
             database_subset: &None,
             only_tables: &vec![],
+            chunk_size: &None,
         };
 
         assert!(p.read(source_options, |_, _| {}).is_ok());
@@ -380,6 +381,7 @@ mod tests {
             skip_config: &vec![],
             database_subset: &None,
             only_tables: &vec![],
+            chunk_size: &None,
         };
 
         assert!(p.read(source_options, |_, _| {}).is_err());
@@ -395,6 +397,7 @@ mod tests {
             skip_config: &vec![],
             database_subset: &None,
             only_tables: &vec![],
+            chunk_size: &None,
         };
 
         p.read(source_options, |original_query, query| {

replibyte/src/source/mysql.rs

@@ -456,6 +456,7 @@ mod tests {
             skip_config: &vec![],
             database_subset: &None,
             only_tables: &vec![],
+            chunk_size: &None,
         };
 
         assert!(p.read(source_options, |_original_query, _query| {}).is_ok());
@@ -468,6 +469,7 @@ mod tests {
             skip_config: &vec![],
             database_subset: &None,
             only_tables: &vec![],
+            chunk_size: &None,
         };
         assert!(p
             .read(source_options, |_original_query, _query| {})
@@ -484,6 +486,7 @@ mod tests {
             skip_config: &vec![],
             database_subset: &None,
             only_tables: &vec![],
+            chunk_size: &None,
         };
         let _ = p.read(source_options, |original_query, query| {
             assert!(original_query.data().len() > 0);

replibyte/src/source/postgres.rs

@@ -581,6 +581,7 @@ mod tests {
             skip_config: &vec![],
             database_subset: &None,
             only_tables: &vec![],
+            chunk_size: &None,
         };
 
         assert!(p.read(source_options, |original_query, query| {}).is_ok());
@@ -593,6 +594,7 @@ mod tests {
             skip_config: &vec![],
             database_subset: &None,
             only_tables: &vec![],
+            chunk_size: &None,
         };
 
         assert!(p.read(source_options, |original_query, query| {}).is_err());
@@ -608,6 +610,7 @@ mod tests {
             skip_config: &vec![],
             database_subset: &None,
             only_tables: &vec![],
+            chunk_size: &None,
         };
 
         let _ = p.read(source_options, |original_query, query| {
@@ -734,6 +737,7 @@ mod tests {
             skip_config: &vec![],
             database_subset: &None,
             only_tables: &vec![],
+            chunk_size: &None,
         };
 
         let _ = p.read(source_options, |original_query, query| {
@@ -775,6 +779,7 @@ mod tests {
             skip_config: &skip_config,
             database_subset: &None,
             only_tables: &vec![],
+            chunk_size: &None,
         };
 
         let _ = p.read(source_options, |_original_query, query| {
@@ -826,6 +831,7 @@ mod tests {
                 passthrough_tables: None,
             }),
             only_tables: &vec![],
+            chunk_size: &None,
         };
 
         let mut rows_percent_50 = vec![];
@@ -863,6 +869,7 @@ mod tests {
                 passthrough_tables: None,
             }),
             only_tables: &vec![],
+            chunk_size: &None,
         };
 
         let mut rows_percent_30 = vec![];

replibyte/src/tasks/full_dump.rs

@@ -69,8 +69,9 @@ where
             Ok(())
         });
 
-        // buffer of 100MB in memory to use and re-use to upload data into datastore
-        let buffer_size = 100 * 1024 * 1024;
+        // buffer default of 100MB (unless specified) in memory to use and re-use to upload data into datastore
+        let chunk_size = self.options.chunk_size.unwrap_or(100);
+        let buffer_size = chunk_size * 1024 * 1024;
         let mut queries = vec![];
         let mut consumed_buffer_size = 0usize;
         let mut total_transferred_bytes = 0usize;