extern crate serde; #[macro_use] extern crate serde_derive; // This allows inserting JSON documents #[macro_use] extern crate serde_json; extern crate ledb; // This allows define typed documents easy #[macro_use] extern crate ledb_derive; extern crate ledb_types; extern crate zip; use clap::{Parser, Subcommand}; use rand::seq::IteratorRandom; use std::path::PathBuf; use std::time::Instant; use std::{fs, io}; use ledb::{Options, Storage}; const ZIP_FILENAME: &str = "json.zip"; const DB_DIR: &str = "db"; macro_rules! make { ($Target:ident; by {$($field:ident),+}; from $src:expr) => {$Target {$( $field: $src.$field ),+}}; ($Target:ident; with defaults and by {$($field:ident),+}; from $src:expr) => {$Target {$( $field: $src.$field ),+ ,..$Target::default()}} } #[derive(Subcommand, Debug)] enum Command { Write, Compact, Print { #[clap(value_parser, default_value = "0")] id: u32, }, ZipPrint { #[clap(value_parser, default_value = "0")] file_num: usize, #[clap(value_parser, default_value = "0")] num: usize, }, } #[derive(Parser, Debug)] #[clap(author, version, about, long_about = None)] #[clap(propagate_version = true)] struct Cli { #[clap(subcommand)] command: Command, #[clap(short, long, action)] measure: bool, } #[derive(Debug, Default, Clone, Serialize, Deserialize)] struct SourceQuestion { #[serde(default)] num: u32, id: String, #[serde(alias = "Вопрос")] description: String, #[serde(alias = "Ответ")] answer: String, #[serde(alias = "Автор")] #[serde(default)] author: String, #[serde(alias = "Комментарий")] #[serde(default)] comment: String, #[serde(alias = "Комментарии")] #[serde(alias = "Инфо")] #[serde(default)] comment1: String, #[serde(alias = "Тур")] #[serde(default)] tour: String, #[serde(alias = "Ссылка")] #[serde(alias = "URL")] #[serde(default)] url: String, #[serde(alias = "Дата")] #[serde(default)] date: String, #[serde(alias = "Обработан")] #[serde(default)] processed_by: String, #[serde(alias = "Редактор")] #[serde(default)] redacted_by: String, #[serde(alias = "Копирайт")] #[serde(default)] copyright: String, #[serde(alias = "Тема")] #[serde(default)] theme: String, #[serde(alias = "Вид")] #[serde(alias = "Тип")] #[serde(default)] kind: String, #[serde(alias = "Источник")] #[serde(default)] source: String, #[serde(alias = "Рейтинг")] #[serde(default)] rating: String, } #[derive(Debug, Default, Clone, Serialize, Deserialize)] struct SourceQuestionsBatch { #[serde(default)] filename: String, #[serde(alias = "Пакет")] #[serde(alias = "Чемпионат")] description: String, #[serde(alias = "Автор")] #[serde(default)] author: String, #[serde(alias = "Комментарий")] #[serde(alias = "Комментарии")] #[serde(alias = "Инфо")] #[serde(default)] comment: String, #[serde(alias = "Ссылка")] #[serde(alias = "URL")] #[serde(default)] url: String, #[serde(alias = "Дата")] #[serde(default)] date: String, #[serde(alias = "Обработан")] #[serde(default)] processed_by: String, #[serde(alias = "Редактор")] #[serde(default)] redacted_by: String, #[serde(alias = "Копирайт")] #[serde(default)] copyright: String, #[serde(alias = "Тема")] #[serde(default)] theme: String, #[serde(alias = "Вид")] #[serde(alias = "Тип")] #[serde(default)] kind: String, #[serde(alias = "Источник")] #[serde(default)] source: String, #[serde(alias = "Рейтинг")] #[serde(default)] rating: String, #[serde(alias = "Вопросы")] questions: Vec, } #[derive(Debug, Default, Clone, Serialize, Deserialize, Document)] struct BatchInfo { #[document(primary)] #[serde(default)] filename: String, #[serde(default)] description: String, #[serde(default)] author: String, #[serde(default)] comment: String, #[serde(default)] url: String, #[serde(default)] date: String, #[serde(default)] processed_by: String, #[serde(default)] redacted_by: String, #[serde(default)] copyright: String, #[serde(default)] theme: String, #[serde(default)] kind: String, #[serde(default)] source: String, #[serde(default)] rating: String, } #[derive(Debug, Default, Clone, Serialize, Deserialize, Document)] struct Question { #[document(primary)] #[serde(default)] num: u32, #[document(index)] id: String, description: String, answer: String, #[serde(default)] author: String, #[serde(default)] comment: String, #[serde(default)] comment1: String, #[serde(default)] tour: String, #[serde(default)] url: String, #[serde(default)] date: String, #[serde(default)] processed_by: String, #[serde(default)] redacted_by: String, #[serde(default)] copyright: String, #[serde(default)] theme: String, #[serde(default)] kind: String, #[serde(default)] source: String, #[serde(default)] rating: String, #[document(nested)] #[serde(default)] batch_info: BatchInfo, } impl From for Question { fn from(src: SourceQuestion) -> Self { make! {Self; with defaults and by { num, id, description, answer, author, comment, comment1, tour, url, date, processed_by, redacted_by, copyright, theme, kind, source, rating }; from src} } } impl From for BatchInfo { fn from(src: SourceQuestionsBatch) -> Self { make! {Self; by { filename, description, author, comment, url, date, processed_by, redacted_by, copyright, theme, kind, source, rating }; from src} } } impl From for Vec { fn from(src: SourceQuestionsBatch) -> Self { let mut result: Vec = src .questions .iter() .map(|item| item.clone().into()) .collect(); let batch_info = BatchInfo::from(src); result.iter_mut().for_each(|mut question| { question.batch_info = batch_info.clone(); }); result } } struct SourceQuestionsZipReader where R: std::io::Read + std::io::Seek, { zipfile: zip::ZipArchive, index: Option, } impl SourceQuestionsZipReader where R: std::io::Read + std::io::Seek, { fn new(zipfile: zip::ZipArchive) -> Self { SourceQuestionsZipReader { zipfile, index: None, } } } impl Iterator for SourceQuestionsZipReader where R: std::io::Read + std::io::Seek, { type Item = (String, Result); fn next(&mut self) -> Option { if self.index.is_none() && !self.zipfile.is_empty() { self.index = Some(0); } match self.index { Some(i) if i < self.zipfile.len() => { self.index = Some(i + 1); self.nth(i) } _ => None, } } fn nth(&mut self, n: usize) -> Option { if self.zipfile.len() <= n { return None; } self.index = Some(n + 1); let file = self.zipfile.by_index(n).unwrap(); let name = file.mangled_name(); let name_str = name.to_str().unwrap(); let data: Result = serde_json::from_reader(file); Some((String::from(name_str), data)) } fn size_hint(&self) -> (usize, Option) { let len = self.zipfile.len(); let index = self.index.unwrap_or(0); let rem = if len > index + 1 { len - (index + 1) } else { 0 }; (rem, Some(rem)) } fn count(self) -> usize where Self: Sized, { self.zipfile.len() } } impl ExactSizeIterator for SourceQuestionsZipReader where R: std::io::Read + std::io::Seek, { fn len(&self) -> usize { self.zipfile.len() } } trait ReadSourceQuestionsBatches where R: std::io::Read + std::io::Seek, { fn source_questions(self) -> SourceQuestionsZipReader; } impl ReadSourceQuestionsBatches for zip::ZipArchive where R: std::io::Read + std::io::Seek, { fn source_questions(self) -> SourceQuestionsZipReader { SourceQuestionsZipReader::new(self) } } fn write_db() { let out_file: PathBuf = [DB_DIR, "data.mdb"].into_iter().collect(); match fs::metadata(&out_file) { Ok(x) if x.is_file() => { fs::remove_file(&out_file).unwrap(); println!(r#""{}" removed"#, out_file.to_str().unwrap()); } _ => {} }; let zip_file = fs::File::open(ZIP_FILENAME).unwrap(); let zip_reader = io::BufReader::new(zip_file); let archive = zip::ZipArchive::new(zip_reader).unwrap(); let options: Options = serde_json::from_value(json!({ "map_size": 900 * 1024 * 1024, // 900mb "write_map": true, "map_async": true, "no_lock": true, "no_meta_sync": true, "no_sync": true, })) .unwrap(); let storage = Storage::new(DB_DIR, options).unwrap(); let collection = storage.collection("questions").unwrap(); println!("converting..."); let mut count: usize = 0; let count = &mut count; archive .source_questions() .filter(|(_, data)| data.is_ok()) .flat_map(|(filename, data)| { let mut data = data.unwrap(); data.filename = filename; let questions: Vec = data.into(); questions }) .for_each(|question| { let result = collection.insert(&question); if result.is_err() { println!("-- {:#?}", question); panic!("{:#?}", result); } else { *count += 1; } }); println!("inserted {}\nwriting...", count); storage.sync(true).unwrap(); print!("stats: "); let stats = storage.stat().unwrap(); println!("{:?}", stats); drop(storage); } fn print_question_from(get_q: F) where F: FnOnce() -> Option, { let q = get_q().unwrap(); println!("{:#?}", q) } fn read_from_zip(file_num: usize, mut num: usize) -> Option { let mut rng = rand::thread_rng(); let zip_file = fs::File::open(ZIP_FILENAME).unwrap(); let zip_reader = io::BufReader::new(zip_file); let archive = zip::ZipArchive::new(zip_reader).unwrap(); let mut source_questions = archive.source_questions(); let (_, batch) = if file_num == 0 { source_questions.choose(&mut rng).unwrap() } else { source_questions.nth(file_num - 1).unwrap() }; let batch = batch.unwrap(); let questions: Vec = batch.into(); if num == 0 { num = (1..=questions.len()).choose(&mut rng).unwrap(); } Some(questions[num - 1].clone()) } fn compact_db() { let options: Options = serde_json::from_value(json!({ "write_map": true, "map_async": true, "no_lock": true, "no_meta_sync": true, "no_sync": true, "compact": true, })) .unwrap(); let storage = Storage::new(DB_DIR, options).unwrap(); storage.sync(true).unwrap(); let stats = storage.stat().unwrap(); println!("{:?}", stats); drop(storage); } fn read_from_db(mut id: u32) -> Option { let options: Options = serde_json::from_value(json!({ "read_only": true, "map_async": true, "no_lock": true, })) .unwrap(); let storage = Storage::new(DB_DIR, options).unwrap(); let collection = storage.collection("questions").unwrap(); let mut rng = rand::thread_rng(); if id == 0 { let last_id = collection.last_id().unwrap(); id = (1..=last_id).choose(&mut rng).unwrap(); } collection.get::(id).unwrap() } // measure and return time elapsed in `func` in seconds pub fn measure(func: F) -> f64 { let start = Instant::now(); func(); let elapsed = start.elapsed(); (elapsed.as_secs() as f64) + (elapsed.subsec_nanos() as f64 / 1_000_000_000.0) } pub fn measure_and_print(func: F) { let m = measure(func); eprintln!("{}", m); } fn main() { let args = Cli::parse(); let mut action: Box = match &args.command { Command::Write => Box::new(write_db), Command::Compact => Box::new(compact_db), Command::Print { id } => { let get_question = Box::new(|| read_from_db(*id)); Box::new(|| print_question_from(get_question)) } Command::ZipPrint { file_num, num } => { let get_question = Box::new(|| read_from_zip(*file_num, *num)); Box::new(|| print_question_from(get_question)) } }; if args.measure { action = Box::new(|| measure_and_print(action)); } action(); }