diff --git a/modules/xml/Cargo.toml b/modules/xml/Cargo.toml index 5a5e812b..f39821f6 100644 --- a/modules/xml/Cargo.toml +++ b/modules/xml/Cargo.toml @@ -31,3 +31,7 @@ structopt = "0.3" color-eyre = "0.5" serde_json = "1.0.61" serde_path_to_error = "0.1" + +[dev-dependencies.rusqlite] +version = "0.26.3" +features = ["bundled", "column_decltype"] diff --git a/modules/xml/examples/add-locale-to-sqlite.rs b/modules/xml/examples/add-locale-to-sqlite.rs new file mode 100644 index 00000000..98aed39c --- /dev/null +++ b/modules/xml/examples/add-locale-to-sqlite.rs @@ -0,0 +1,212 @@ +use std::{ + collections::{HashMap, HashSet}, + fs::File, + io::BufReader, + path::PathBuf, + time::Instant, +}; + +use argh::FromArgs; +use color_eyre::eyre::WrapErr; +use quick_xml::{events::Event as XmlEvent, Reader as XmlReader}; +use rusqlite::Connection; + +use assembly_xml::common::exact::{ + expect_attribute, expect_child_or_end, expect_end, expect_start, expect_text, + expect_text_or_end, Error, +}; + +const TAG_LOCALIZATION: &str = "localization"; +const TAG_LOCALES: &str = "locales"; +const TAG_LOCALE: &str = "locale"; +const TAG_PHRASES: &str = "phrases"; +const TAG_PHRASE: &str = "phrase"; +const TAG_TRANSLATION: &str = "translation"; + +const ATTR_COUNT: &str = "count"; +const ATTR_LOCALE: &str = "locale"; +const ATTR_ID: &str = "id"; + +/// Try to add locale info to an existing converted cdclient SQLite file +/// +/// This function does the following: +/// +/// 1. `BEGIN`s a transaction +/// 2. For every locale entry matching a specific format: +/// a. If not already done, run ALTER TABLE ADD COLUMN to add the column to the DB +/// b. Runs UPDATE to add the info to the DB +/// 3. `COMMIT`s the transaction +pub fn try_add_locale( + conn: &mut Connection, + mut reader: XmlReader>, +) -> Result<(), Error> { + // All data we will be inserting will be strings, so disable the check for string type for better performance + conn.execute("PRAGMA ignore_check_constraints=ON", rusqlite::params![]) + .unwrap(); + conn.execute("BEGIN", rusqlite::params![]).unwrap(); + + let mut buf = Vec::new(); + + // The `Reader` does not implement `Iterator` because it outputs borrowed data (`Cow`s) + if let Ok(XmlEvent::Decl(_)) = reader.read_event(&mut buf) {} + buf.clear(); + + let _ = expect_start(TAG_LOCALIZATION, &mut reader, &mut buf)?; + buf.clear(); + + let e_locales = expect_start(TAG_LOCALES, &mut reader, &mut buf)?; + let locale_count: usize = expect_attribute(ATTR_COUNT, &reader, &e_locales)?; + let mut real_locale_count = 0; + buf.clear(); + + while expect_child_or_end(TAG_LOCALE, TAG_LOCALES, &mut reader, &mut buf)?.is_some() { + buf.clear(); + + let _locale = expect_text(&mut reader, &mut buf)?; + + expect_end(TAG_LOCALE, &mut reader, &mut buf)?; + buf.clear(); + + real_locale_count += 1; + } + buf.clear(); + + if real_locale_count != locale_count { + println!( + "locale.xml specifies a locale count of {}, but has {}", + locale_count, real_locale_count + ); + } + + let e_phrases = expect_start(TAG_PHRASES, &mut reader, &mut buf)?; + let phrase_count: usize = expect_attribute(ATTR_COUNT, &reader, &e_phrases)?; + let mut real_phrase_count = 0; + buf.clear(); + + let mut tables: HashMap> = HashMap::new(); + let mut table_pk: HashMap = HashMap::new(); + + let mut stmt = conn.prepare("SELECT name, sql from sqlite_master").unwrap(); + let mut rows = stmt.query([]).unwrap(); + while let Some(row) = rows.next().unwrap() { + let name = row.get::(0).unwrap(); + let sql = row.get::(1).unwrap(); + let start = match sql.find('[') { + None => continue, + Some(x) => x, + }; + let end = match sql.find(']') { + None => continue, + Some(x) => x, + }; + let pk = &sql[start + 1..end]; + table_pk.insert(name, pk.to_string()); + } + + while let Some(e_phrase) = expect_child_or_end(TAG_PHRASE, TAG_PHRASES, &mut reader, &mut buf)? + { + let id: String = expect_attribute(ATTR_ID, &reader, &e_phrase)?; + buf.clear(); + + while let Some(e_translation) = + expect_child_or_end(TAG_TRANSLATION, TAG_PHRASE, &mut reader, &mut buf)? + { + let locale: String = expect_attribute(ATTR_LOCALE, &reader, &e_translation)?; + buf.clear(); + + let trans = expect_text_or_end(TAG_TRANSLATION, &mut reader, &mut buf)?; + + let mut splitn = id.splitn(3, '_'); + let table = match splitn.next() { + Some(x) => String::from(x), + None => continue, + }; + if table.chars().all(|c| !c.is_ascii_lowercase()) + || table.chars().all(|c| !c.is_ascii_uppercase()) + { + continue; + } + let pk = match splitn.next() { + Some(x) => x, + None => continue, + }; + let column = match splitn.next() { + Some(x) => x, + None => continue, + }; + + let columns = match tables.get_mut(&table) { + Some(x) => x, + None => { + tables.insert(table.to_string(), HashSet::new()); + tables.get_mut(&table).unwrap() + } + }; + let col_loc = format!("{}_{}", column, locale); + let update = format!( + "UPDATE \"{}\" SET \"{}\" = ? WHERE \"{}\" = ?", + table, + col_loc, + table_pk.get(&table).unwrap() + ); + if !columns.contains(&col_loc) { + let sql = format!("ALTER TABLE \"{}\" ADD COLUMN \"{}\" TEXT4 CHECK (TYPEOF(\"{}\") in ('text', 'null'))", table, col_loc, col_loc); + conn.execute(&sql, rusqlite::params![]).unwrap(); + columns.insert(col_loc); + } + conn.execute(&update, rusqlite::params![trans, pk]).unwrap(); + buf.clear(); + } + buf.clear(); + + real_phrase_count += 1; + } + buf.clear(); + + if phrase_count != real_phrase_count { + println!( + "locale.xml specifies a count of {} phrases, but has {}", + phrase_count, real_phrase_count + ); + } + conn.execute("COMMIT", rusqlite::params![]).unwrap(); + Ok(()) +} + +#[derive(FromArgs)] +/// Turns an FDB file into an equivalent SQLite file +struct Options { + /// the locale.xml file + #[argh(positional)] + locale: PathBuf, + /// the SQLite DB to augment + #[argh(positional)] + sqlite: PathBuf, +} + +fn main() -> color_eyre::Result<()> { + color_eyre::install()?; + let opts: Options = argh::from_env(); + let start = Instant::now(); + + println!("Copying data, this may take a few seconds..."); + + let mut conn = Connection::open(opts.sqlite)?; + + let file = File::open(&opts.locale)?; + let file = BufReader::new(file); + + let mut reader = XmlReader::from_reader(file); + reader.trim_text(true); + + try_add_locale(&mut conn, reader).wrap_err("Failed to add locale info to sqlite")?; + + let duration = start.elapsed(); + println!( + "Finished in {}.{}s", + duration.as_secs(), + duration.subsec_millis() + ); + + Ok(()) +}