diff --git a/.gitignore b/.gitignore index 2cef3db..cc7f9bb 100644 --- a/.gitignore +++ b/.gitignore @@ -2,3 +2,4 @@ **/*.rs.bk *.data *.old +out.sql \ No newline at end of file diff --git a/src/main.rs b/src/main.rs index 9f276da..8bc9e51 100644 --- a/src/main.rs +++ b/src/main.rs @@ -92,9 +92,19 @@ impl FromStr for LevelSizes { } } -fn main() { - #[allow(deprecated)] - let matches = App::new(crate_name!()) +struct Config { + db_url: String, + output_file: Option, + room_id: String, + max_state_group: Option, + min_saved_rows: Option, + transactions: bool, + level_sizes: LevelSizes, +} + +impl Config { + fn parse_arguments() -> Config { + let matches = App::new(crate_name!()) .version(crate_version!()) .author(crate_authors!("\n")) .about(crate_description!()) @@ -156,33 +166,50 @@ fn main() { .takes_value(true), ).get_matches(); - let db_url = matches - .value_of("postgres-url") - .expect("db url should be required"); + let db_url = matches + .value_of("postgres-url") + .expect("db url should be required"); - let mut output_file = matches - .value_of("output_file") - .map(|path| File::create(path).unwrap()); + let output_file = matches + .value_of("output_file") + .map(|path| File::create(path).unwrap()); - let room_id = matches - .value_of("room_id") - .expect("room_id should be required since no file"); + let room_id = matches + .value_of("room_id") + .expect("room_id should be required since no file"); - let max_state_group = matches - .value_of("max_state_group") - .map(|s| s.parse().expect("max_state_group must be an integer")); + let max_state_group = matches + .value_of("max_state_group") + .map(|s| s.parse().expect("max_state_group must be an integer")); - let min_saved_rows = matches - .value_of("min_saved_rows") - .map(|v| v.parse().expect("COUNT must be an integer")); + let min_saved_rows = matches + .value_of("min_saved_rows") + .map(|v| v.parse().expect("COUNT must be an integer")); - let transactions = matches.is_present("transactions"); + let transactions = matches.is_present("transactions"); - let level_sizes = value_t_or_exit!(matches, "level_sizes", LevelSizes); + let level_sizes = value_t_or_exit!(matches, "level_sizes", LevelSizes); + + Config { + db_url: String::from(db_url), + output_file, + room_id: String::from(room_id), + max_state_group, + min_saved_rows, + transactions, + level_sizes, + } + } +} + +fn main() { + let mut config = Config::parse_arguments(); // First we need to get the current state groups - println!("Fetching state from DB for room '{}'...", room_id); - let state_group_map = database::get_data_from_db(db_url, room_id, max_state_group); + println!("Fetching state from DB for room '{}'...", config.room_id); + + let state_group_map = + database::get_data_from_db(&config.db_url, &config.room_id, config.max_state_group); println!("Number of state groups: {}", state_group_map.len()); @@ -196,7 +223,7 @@ fn main() { println!("Compressing state..."); - let compressor = Compressor::compress(&state_group_map, &level_sizes.0); + let compressor = Compressor::compress(&state_group_map, &config.level_sizes.0); let new_state_group_map = compressor.new_state_group_map; @@ -228,7 +255,7 @@ fn main() { compressor.stats.state_groups_changed ); - if let Some(min) = min_saved_rows { + if let Some(min) = config.min_saved_rows { let saving = (original_summed_size - compressed_summed_size) as i32; if saving < min { println!( @@ -239,25 +266,39 @@ fn main() { } } + check_that_maps_match(&state_group_map, &new_state_group_map); + // If we are given an output file, we output the changes as SQL. If the // `transactions` argument is set we wrap each change to a state group in a // transaction. - if let Some(output) = &mut output_file { - println!("Writing changes..."); + output_sql(&mut config, &state_group_map, &new_state_group_map); +} - let pb = ProgressBar::new(state_group_map.len() as u64); - pb.set_style( - ProgressStyle::default_bar().template("[{elapsed_precise}] {bar} {pos}/{len} {msg}"), - ); - pb.set_message("state groups"); - pb.enable_steady_tick(100); +fn output_sql( + config: &mut Config, + old_map: &BTreeMap, + new_map: &BTreeMap, +) { + if config.output_file.is_none() { + return; + } - for (sg, old_entry) in &state_group_map { - let new_entry = &new_state_group_map[sg]; + println!("Writing changes..."); + + let pb = ProgressBar::new(old_map.len() as u64); + pb.set_style( + ProgressStyle::default_bar().template("[{elapsed_precise}] {bar} {pos}/{len} {msg}"), + ); + pb.set_message("state groups"); + pb.enable_steady_tick(100); + + if let Some(output) = &mut config.output_file { + for (sg, old_entry) in old_map { + let new_entry = &new_map[sg]; if old_entry != new_entry { - if transactions { + if config.transactions { writeln!(output, "BEGIN;").unwrap(); } @@ -292,7 +333,7 @@ fn main() { output, "({}, {}, {}, {}, {})", sg, - PGEscape(room_id), + PGEscape(&config.room_id), PGEscape(t), PGEscape(s), PGEscape(e) @@ -302,7 +343,7 @@ fn main() { writeln!(output, ";").unwrap(); } - if transactions { + if config.transactions { writeln!(output, "COMMIT;").unwrap(); } writeln!(output).unwrap(); @@ -310,13 +351,18 @@ fn main() { pb.inc(1); } - - pb.finish(); } + pb.finish(); +} + +fn check_that_maps_match( + old_map: &BTreeMap, + new_map: &BTreeMap, +) { println!("Checking that state maps match..."); - let pb = ProgressBar::new(state_group_map.len() as u64); + let pb = ProgressBar::new(old_map.len() as u64); pb.set_style( ProgressStyle::default_bar().template("[{elapsed_precise}] {bar} {pos}/{len} {msg}"), ); @@ -325,11 +371,11 @@ fn main() { // Now let's iterate through and assert that the state for each group // matches between the two versions. - state_group_map + old_map .par_iter() // This uses rayon to run the checks in parallel .try_for_each(|(sg, _)| { - let expected = collapse_state_maps(&state_group_map, *sg); - let actual = collapse_state_maps(&new_state_group_map, *sg); + let expected = collapse_state_maps(&old_map, *sg); + let actual = collapse_state_maps(&new_map, *sg); pb.inc(1);