Split up main (#39)

Moved code from inside the main function into function calls
This commit is contained in:
Azrenbeth
2021-07-27 15:31:48 +01:00
committed by GitHub
parent f3b1dd0b8a
commit 31e37996ea
2 changed files with 90 additions and 43 deletions

1
.gitignore vendored
View File

@@ -2,3 +2,4 @@
**/*.rs.bk **/*.rs.bk
*.data *.data
*.old *.old
out.sql

View File

@@ -92,8 +92,18 @@ impl FromStr for LevelSizes {
} }
} }
fn main() { struct Config {
#[allow(deprecated)] db_url: String,
output_file: Option<File>,
room_id: String,
max_state_group: Option<i64>,
min_saved_rows: Option<i32>,
transactions: bool,
level_sizes: LevelSizes,
}
impl Config {
fn parse_arguments() -> Config {
let matches = App::new(crate_name!()) let matches = App::new(crate_name!())
.version(crate_version!()) .version(crate_version!())
.author(crate_authors!("\n")) .author(crate_authors!("\n"))
@@ -160,7 +170,7 @@ fn main() {
.value_of("postgres-url") .value_of("postgres-url")
.expect("db url should be required"); .expect("db url should be required");
let mut output_file = matches let output_file = matches
.value_of("output_file") .value_of("output_file")
.map(|path| File::create(path).unwrap()); .map(|path| File::create(path).unwrap());
@@ -180,9 +190,26 @@ fn main() {
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 // First we need to get the current state groups
println!("Fetching state from DB for room '{}'...", room_id); println!("Fetching state from DB for room '{}'...", config.room_id);
let state_group_map = database::get_data_from_db(db_url, room_id, max_state_group);
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()); println!("Number of state groups: {}", state_group_map.len());
@@ -196,7 +223,7 @@ fn main() {
println!("Compressing state..."); 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; let new_state_group_map = compressor.new_state_group_map;
@@ -228,7 +255,7 @@ fn main() {
compressor.stats.state_groups_changed 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; let saving = (original_summed_size - compressed_summed_size) as i32;
if saving < min { if saving < min {
println!( 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 // 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 // `transactions` argument is set we wrap each change to a state group in a
// transaction. // transaction.
if let Some(output) = &mut output_file { output_sql(&mut config, &state_group_map, &new_state_group_map);
}
fn output_sql(
config: &mut Config,
old_map: &BTreeMap<i64, StateGroupEntry>,
new_map: &BTreeMap<i64, StateGroupEntry>,
) {
if config.output_file.is_none() {
return;
}
println!("Writing changes..."); println!("Writing changes...");
let pb = ProgressBar::new(state_group_map.len() as u64); let pb = ProgressBar::new(old_map.len() as u64);
pb.set_style( pb.set_style(
ProgressStyle::default_bar().template("[{elapsed_precise}] {bar} {pos}/{len} {msg}"), ProgressStyle::default_bar().template("[{elapsed_precise}] {bar} {pos}/{len} {msg}"),
); );
pb.set_message("state groups"); pb.set_message("state groups");
pb.enable_steady_tick(100); pb.enable_steady_tick(100);
for (sg, old_entry) in &state_group_map { if let Some(output) = &mut config.output_file {
let new_entry = &new_state_group_map[sg]; for (sg, old_entry) in old_map {
let new_entry = &new_map[sg];
if old_entry != new_entry { if old_entry != new_entry {
if transactions { if config.transactions {
writeln!(output, "BEGIN;").unwrap(); writeln!(output, "BEGIN;").unwrap();
} }
@@ -292,7 +333,7 @@ fn main() {
output, output,
"({}, {}, {}, {}, {})", "({}, {}, {}, {}, {})",
sg, sg,
PGEscape(room_id), PGEscape(&config.room_id),
PGEscape(t), PGEscape(t),
PGEscape(s), PGEscape(s),
PGEscape(e) PGEscape(e)
@@ -302,7 +343,7 @@ fn main() {
writeln!(output, ";").unwrap(); writeln!(output, ";").unwrap();
} }
if transactions { if config.transactions {
writeln!(output, "COMMIT;").unwrap(); writeln!(output, "COMMIT;").unwrap();
} }
writeln!(output).unwrap(); writeln!(output).unwrap();
@@ -310,13 +351,18 @@ fn main() {
pb.inc(1); pb.inc(1);
} }
}
pb.finish(); pb.finish();
} }
fn check_that_maps_match(
old_map: &BTreeMap<i64, StateGroupEntry>,
new_map: &BTreeMap<i64, StateGroupEntry>,
) {
println!("Checking that state maps match..."); 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( pb.set_style(
ProgressStyle::default_bar().template("[{elapsed_precise}] {bar} {pos}/{len} {msg}"), 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 // Now let's iterate through and assert that the state for each group
// matches between the two versions. // matches between the two versions.
state_group_map old_map
.par_iter() // This uses rayon to run the checks in parallel .par_iter() // This uses rayon to run the checks in parallel
.try_for_each(|(sg, _)| { .try_for_each(|(sg, _)| {
let expected = collapse_state_maps(&state_group_map, *sg); let expected = collapse_state_maps(&old_map, *sg);
let actual = collapse_state_maps(&new_state_group_map, *sg); let actual = collapse_state_maps(&new_map, *sg);
pb.inc(1); pb.inc(1);