rsbag/examples/mongobag.rs
2021-11-22 00:07:58 -08:00

66 lines
1.9 KiB
Rust

use std::env::args;
use bson::Document;
use indicatif::ParallelProgressIterator;
use log::info;
use mongodb::sync::Client;
use rayon::iter::ParallelIterator;
use rsbag::{chunk::ChunkMessageIterator, Bag, Result};
fn main() -> Result<()> {
color_eyre::install()?;
env_logger::init();
let args: Vec<_> = args().collect();
if args.len() != 2 {
eprintln!("Usage: {} <bag path>", args[0]);
return Ok(());
}
let bag_path = &args[1];
let mut bag = Bag::open(bag_path)?;
let layouts = bag.compute_message_layouts()?;
let info = bag.compute_info()?;
let total_messages: u64 = info.per_connection.values().sum();
info!("exporting {} messages", total_messages);
let client = Client::with_uri_str("mongodb://localhost:27017")?;
let database = client.database("rsbag");
let collection = database.collection::<Document>("messages");
let num_chunks = bag.index().chunks.len();
bag.read_chunks()
.progress_count(num_chunks as u64)
.for_each_with(collection, |collection, chunk| {
let it = ChunkMessageIterator::new(chunk.unwrap());
let mut documents = Vec::new();
for msg in it {
let msg = msg.unwrap();
let (topic, msg) = layouts.decode_bson(&msg).unwrap();
documents.push(bson::doc! {
"topic": topic,
"msg": msg,
});
}
collection.insert_many(documents, None).unwrap();
});
// bag.read_messages()
// .progress_count(total_messages)
// .map(|msg| {
// let msg = msg.unwrap();
// layouts.decode_bson(&msg).unwrap()
// })
// .for_each_with(collection, |collection, (topic, msg)| {
// collection.insert_one(bson::doc! {
// "topic": topic,
// "msg": msg,
// }, None).unwrap();
// });
Ok(())
}