diff options
Diffstat (limited to 'rust/src')
-rw-r--r-- | rust/src/bin/fatcat-export.rs | 161 | ||||
-rw-r--r-- | rust/src/lib.rs | 2 |
2 files changed, 162 insertions, 1 deletions
diff --git a/rust/src/bin/fatcat-export.rs b/rust/src/bin/fatcat-export.rs new file mode 100644 index 00000000..d7371c87 --- /dev/null +++ b/rust/src/bin/fatcat-export.rs @@ -0,0 +1,161 @@ +//! JSON Export Helper + +extern crate clap; +extern crate diesel; +extern crate dotenv; +#[macro_use] +extern crate error_chain; +extern crate fatcat; +extern crate fatcat_api_spec; +#[macro_use] +extern crate slog; +extern crate slog_async; +extern crate slog_term; +extern crate uuid; +extern crate crossbeam_channel; +extern crate serde_json; + +use clap::{App, Arg}; +use slog::{Drain, Logger}; +use dotenv::dotenv; +use std::env; + +use diesel::prelude::*; +use diesel::r2d2::ConnectionManager; +use fatcat_api_spec::models::*; +use std::str::FromStr; +use uuid::Uuid; +use fatcat::ConnectionPool; +use fatcat::api_helpers::*; +use fatcat::api_entity_crud::*; +use fatcat::errors::*; + +use std::thread; +//use std::io::{Stdout,StdoutLock}; +use crossbeam_channel as channel; +//use num_cpus; TODO: +use std::io::prelude::*; +use std::io::{BufReader, BufWriter}; + +const CHANNEL_BUFFER_LEN: usize = 200; + +struct IdentRow { + ident_id: FatCatId, + rev_id: Option<Uuid>, + redirect_id: Option<FatCatId>, +} + +/// Instantiate a new API server with a pooled database connection +pub fn database_worker_pool() -> Result<ConnectionPool> { + dotenv().ok(); + let database_url = env::var("DATABASE_URL").expect("DATABASE_URL must be set"); + let manager = ConnectionManager::<PgConnection>::new(database_url); + let pool = diesel::r2d2::Pool::builder() + .build(manager) + .expect("Failed to create database pool."); + Ok(pool) +} + +fn loop_work(row_receiver: channel::Receiver<IdentRow>, output_sender: channel::Sender<String>, db_conn: &DbConn, expand: Option<ExpandFlags>) -> Result<()> { + for row in row_receiver { + // TODO: based on types + let mut entity = ReleaseEntity::db_get_rev(db_conn, row.rev_id.expect("valid, non-deleted row"))?; + entity.state = Some("active".to_string()); // TODO + entity.ident = Some(row.ident_id.to_string()); + if let Some(expand) = expand { + entity.db_expand(db_conn, expand)?; + } + output_sender.send(serde_json::to_string(&entity)?); + } + Ok(()) +} + +fn loop_printer(output_receiver: channel::Receiver<String>, done_sender: channel::Sender<()>) -> Result<()> { + let output = std::io::stdout(); + // XXX should log... + // let mut buf_output = BufWriter::new(output.lock()); + let mut buf_output = BufWriter::new(output); + for line in output_receiver { + buf_output.write_all(&line.into_bytes())?; + buf_output.write(b"\n")?; + buf_output.flush()?; + } + drop(done_sender); + Ok(()) +} + +fn parse_line(s: String) -> Result<IdentRow> { + let fields: Vec<String> = s.split("\t").map(|v| v.to_string()).collect(); + if fields.len() != 3 { + bail!("Invalid input line"); + } + Ok(IdentRow { + ident_id: FatCatId::from_uuid(&Uuid::from_str(&fields[0])?), + rev_id: Some(Uuid::from_str(&fields[1])?), + redirect_id: None, + }) +} + +#[test] +fn test_parse_line() { + assert!(false) +} + +// Use num_cpus/2, or CLI arg for worker count +// +// 1. open buffered reader, buffered writer, and database pool. create channels +// 2. spawn workers: +// - get a database connection from database pool +// - iterate over row channel, pushing Strings to output channel +// - exit when end of rows +// 3. spawn output printer +// 4. read rows, pushing to row channel +// => every N lines, log to stderr +// 5. wait for all channels to finish, and quit +pub fn run() -> Result<()> { + + let db_pool = database_worker_pool()?; + let buf_input = BufReader::new(std::io::stdin()); + let worker_count = 4; + let (row_sender, row_receiver) = channel::bounded(CHANNEL_BUFFER_LEN); + let (output_sender, output_receiver) = channel::bounded(CHANNEL_BUFFER_LEN); + let (done_sender, done_receiver) = channel::bounded(0); + + // Start row worker threads + for _ in 0..worker_count { + let db_conn = db_pool.get().expect("database pool"); + let row_receiver = row_receiver.clone(); + let output_sender = output_sender.clone(); + thread::spawn(move || loop_work(row_receiver, output_sender, &db_conn, None)); + } + drop(output_sender); + // Start printer thread + thread::spawn(move || loop_printer(output_receiver, done_sender)); + + for line in buf_input.lines() { + let line = line?; + let row = parse_line(line)?; + row_sender.send(row); + } + drop(row_sender); + done_receiver.recv(); + Ok(()) +} + +fn main() { + /* + let matches = App::new("server") + .arg( + Arg::with_name("https") + .long("https") + .help("Whether to use HTTPS or not"), + ) + .get_matches(); + + let decorator = slog_term::TermDecorator::new().build(); + let drain = slog_term::CompactFormat::new(decorator).build().fuse(); + let drain = slog_async::Async::new(drain).build().fuse(); + let logger = Logger::root(drain, o!()); + */ + run().expect("success") +} diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 381c67ea..d3ef4a9c 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -29,7 +29,7 @@ pub mod api_wrappers; pub mod database_models; pub mod database_schema; -mod errors { +pub mod errors { // Create the Error, ErrorKind, ResultExt, and Result types error_chain! { foreign_links { Fmt(::std::fmt::Error); |