Add option to commit changes to the database automatically (#53)

src/database.rs

@@ -19,6 +19,8 @@ use postgres_openssl::MakeTlsConnector;
 use rand::{distributions::Alphanumeric, thread_rng, Rng};
 use std::{borrow::Cow, collections::BTreeMap, fmt};
 
+use crate::{generate_sql, Config};
+
 use super::StateGroupEntry;
 
 /// Fetch the entries in state_groups_state (and their prev groups) for a
@@ -339,3 +341,51 @@ fn test_pg_escape() {
     assert_eq!(&s[0..1], "$");
     assert_eq!(&s[start_pos - 1..start_pos], "$");
 }
+
+/// Note that currently ignores config.transactions and wraps every state
+/// group in it's own transaction (i.e. as if config.transactions was true)
+///
+/// # Arguments
+///
+/// * `config` -    A Config struct that contains information
+///                 about the run (e.g. room_id and database url)
+/// * `old_map` -   The state group data originally in the database
+/// * `new_map` -   The state group data generated by the compressor to
+///                 replace replace the old contents
+pub fn send_changes_to_db(
+    config: &Config,
+    old_map: &BTreeMap<i64, StateGroupEntry>,
+    new_map: &BTreeMap<i64, StateGroupEntry>,
+) {
+    // connect to the database
+    let mut builder = SslConnector::builder(SslMethod::tls()).unwrap();
+    builder.set_verify(SslVerifyMode::NONE);
+    let connector = MakeTlsConnector::new(builder.build());
+
+    let mut client = Client::connect(&config.db_url, connector).unwrap();
+
+    println!("Writing changes...");
+
+    // setup the progress bar
+    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);
+
+    for sql_transaction in generate_sql(old_map, new_map, &config.room_id) {
+        // commit this change to the database
+        // N.B. this is a synchronous library so will wait until finished before continueing...
+        // if want to speed up compressor then this might be a good place to start!
+        let mut single_group_transaction = client.transaction().unwrap();
+        single_group_transaction
+            .batch_execute(&sql_transaction)
+            .unwrap();
+        single_group_transaction.commit().unwrap();
+
+        pb.inc(1);
+    }
+
+    pb.finish();
+}