diff --git a/Cargo.lock b/Cargo.lock index 36b2494..2c02249 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -725,6 +725,7 @@ dependencies = [ "fallible-iterator 0.1.5 (registry+https://github.com/rust-lang/crates.io-index)", "indicatif 0.9.0 (registry+https://github.com/rust-lang/crates.io-index)", "postgres 0.15.2 (registry+https://github.com/rust-lang/crates.io-index)", + "rand 0.5.5 (registry+https://github.com/rust-lang/crates.io-index)", "rayon 1.0.2 (registry+https://github.com/rust-lang/crates.io-index)", "state-map 0.1.0 (git+https://github.com/matrix-org/rust-matrix-state-map)", "string_cache 0.7.3 (registry+https://github.com/rust-lang/crates.io-index)", diff --git a/Cargo.toml b/Cargo.toml index dac3c76..3ff9b82 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,6 +9,7 @@ clap = "2.32.0" fallible-iterator = "0.1.5" indicatif = "0.9.0" postgres = "0.15.2" +rand = "0.5.5" rayon = "1.0.2" string_cache = "0.7.3" diff --git a/src/compressor.rs b/src/compressor.rs index ae4674c..7a7279c 100644 --- a/src/compressor.rs +++ b/src/compressor.rs @@ -264,7 +264,8 @@ fn test_new_map() { (10, 9), (11, 10), (13, 12), - ].into_iter() + ] + .into_iter() .collect(); for sg in 0i64..=13i64 { diff --git a/src/database.rs b/src/database.rs index 4efcbee..b46bb3c 100644 --- a/src/database.rs +++ b/src/database.rs @@ -15,8 +15,12 @@ use fallible_iterator::FallibleIterator; use indicatif::{ProgressBar, ProgressStyle}; use postgres::{Connection, TlsMode}; +use rand::distributions::Alphanumeric; +use rand::{thread_rng, Rng}; +use std::borrow::Cow; use std::collections::BTreeMap; +use std::fmt; use StateGroupEntry; @@ -47,7 +51,8 @@ pub fn get_data_from_db(db_url: &str, room_id: &str) -> BTreeMap BTreeMap BTreeMap BTreeMap(pub &'a str); + +impl<'a> fmt::Display for PGEscapse<'a> { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + let mut delim = Cow::from("$$"); + while self.0.contains(&delim as &str) { + let s: String = thread_rng().sample_iter(&Alphanumeric).take(10).collect(); + + delim = format!("${}$", s).into(); + } + + write!(f, "{}{}{}", delim, self.0, delim) + } +} + +#[test] +fn test_pg_escape() { + let s = format!("{}", PGEscapse("test")); + assert_eq!(s, "$$test$$"); + + let dodgy_string = "test$$ing"; + + let s = format!("{}", PGEscapse(dodgy_string)); + + // prefix and suffixes should match + let start_pos = s.find(dodgy_string).expect("expected to find dodgy string"); + let end_pos = start_pos + dodgy_string.len(); + assert_eq!(s[..start_pos], s[end_pos..]); + + // .. and they should start and end with '$' + assert_eq!(&s[0..1], "$"); + assert_eq!(&s[start_pos - 1..start_pos], "$"); +} diff --git a/src/main.rs b/src/main.rs index 58fe243..915d6ac 100644 --- a/src/main.rs +++ b/src/main.rs @@ -12,7 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. - //! This is a tool that attempts to further compress state maps within a //! Synapse instance's database. Specifically, it aims to reduce the number of //! rows that a given room takes up in the `state_groups_state` table. @@ -22,6 +21,7 @@ extern crate clap; extern crate fallible_iterator; extern crate indicatif; extern crate postgres; +extern crate rand; extern crate rayon; extern crate state_map; extern crate string_cache; @@ -30,6 +30,7 @@ mod compressor; mod database; use compressor::Compressor; +use database::PGEscapse; use clap::{App, Arg}; use indicatif::{ProgressBar, ProgressStyle}; @@ -259,7 +260,15 @@ fn main() { } else { write!(output, " ,"); } - writeln!(output, "({}, '{}', '{}', '{}', '{}')", sg, room_id, t, s, e); + writeln!( + output, + "({}, {}, {}, {}, {})", + sg, + PGEscapse(room_id), + PGEscapse(t), + PGEscapse(s), + PGEscapse(e) + ); } writeln!(output, ";"); } @@ -268,7 +277,6 @@ fn main() { writeln!(output, "COMMIT;"); } writeln!(output); - } pb.inc(1); @@ -304,7 +312,8 @@ fn main() { } else { Ok(()) } - }).expect("expected state to match"); + }) + .expect("expected state to match"); pb.finish();