use anyhow::{Context, Result, anyhow}; use gethostname::gethostname; use is_terminal::IsTerminal; use std::collections::HashMap; use std::io::{self, Read, Write}; use clap::Command; use clap::error::ErrorKind; use log::debug; use rusqlite::Connection; use std::path::PathBuf; use crate::compression_engine::get_compression_engine; use crate::db::{self}; use crate::digest_engine::get_digest_engine; use crate::meta_plugin::{MetaPlugin, MetaPluginType, get_meta_plugin}; use crate::modes::common::{cmd_args_compression_type, cmd_args_digest_type, cmd_args_meta_plugin_types, store_item_digest_value, get_meta_from_env, store_item_meta_value}; use chrono::Utc; pub fn mode_save( cmd: &mut Command, args: &crate::Args, ids: &mut Vec, tags: &mut Vec, conn: &mut Connection, data_path: PathBuf, ) -> Result<()> { if !ids.is_empty() { cmd.error( ErrorKind::InvalidValue, "ID given, you cannot supply IDs when using --save", ) .exit(); } if tags.is_empty() { tags.push("none".to_string()); } let digest_type = cmd_args_digest_type(cmd, &args); debug!("MAIN: Digest type: {:?}", digest_type); let mut digest_engine = get_digest_engine(digest_type.clone()); let compression_type = cmd_args_compression_type(cmd, &args); debug!("MAIN: Compression type: {:?}", compression_type); let compression_engine = get_compression_engine(compression_type.clone()).expect("Unable to get compression engine"); let meta_plugin_types: Vec = cmd_args_meta_plugin_types(cmd, &args); debug!("MAIN: Meta plugin types: {:?}", meta_plugin_types); // Initialize meta_plugins with MetaPlugin instances for each MetaPluginType let mut meta_plugins: Vec> = meta_plugin_types .into_iter() .map(|meta_plugin_type| get_meta_plugin(meta_plugin_type)) .collect(); let mut item = db::Item { id: None, ts: Utc::now(), size: None, compression: compression_type.to_string(), }; let id = db::insert_item(conn, item.clone())?; item.id = Some(id); debug!("MAIN: Added item {:?}", item.clone()); if !args.options.quiet { if std::io::stderr().is_terminal() { let mut t = term::stderr().unwrap(); t.reset().unwrap_or(()); t.attr(term::Attr::Bold).unwrap_or(()); write!(t, "KEEP:").unwrap_or(()); t.reset().unwrap_or(()); write!(t, " New item ").unwrap_or(()); t.attr(term::Attr::Bold).unwrap_or(()); write!(t, "{id}")?; t.reset().unwrap_or(()); write!(t, " tags: ")?; t.attr(term::Attr::Bold).unwrap_or(()); write!(t, "{}", tags.join(" "))?; t.reset().unwrap_or(()); writeln!(t)?; std::io::stderr().flush()?; } else { let mut t = std::io::stderr(); writeln!(t, "KEEP: New item: {} tags: {:?}", id, tags)?; } } db::set_item_tags(conn, item.clone(), tags)?; let mut item_meta: HashMap = get_meta_from_env(); if let Ok(hostname) = gethostname().into_string() { if !item_meta.contains_key("hostname") { item_meta.insert("hostname".to_string(), hostname); } } for item in args.item.meta.iter() { let item = item.clone(); item_meta.insert(item.key, item.value); } for kv in item_meta.iter() { let meta = db::Meta { id: item.id.unwrap(), name: kv.0.to_string(), value: kv.1.to_string(), }; db::store_meta(conn, meta)?; } let mut item_path = data_path.clone(); item_path.push(id.to_string()); let mut stdin = io::stdin().lock(); let mut stdout = io::stdout().lock(); let mut buffer = [0; libc::BUFSIZ as usize]; let mut item_out: Box = compression_engine .create(item_path.clone()) .context(anyhow!( "Unable to write file {:?} using compression {:?}", item_path, compression_type ))?; debug!("MAIN: Starting IO loop"); loop { let n = stdin.read(&mut buffer[..libc::BUFSIZ as usize])?; item.size = match item.size { None => Some(n as i64), Some(prev_n) => Some(prev_n + n as i64), }; if n == 0 { debug!("MAIN: EOF on STDIN"); break; } debug!("MAIN: Loop - {:?} bytes", item.size); stdout.write_all(&buffer[..n])?; item_out.write_all(&buffer[..n])?; digest_engine.update(&buffer[..n]); for meta_plugin in meta_plugins.iter_mut() { meta_plugin.update(&buffer[..n]); } } debug!("MAIN: Ending IO loop after {:?} bytes", item.size); stdout.flush()?; item_out.flush()?; let digest = digest_engine.finalize()?; debug!("DIGEST: {}", digest); store_item_digest_value(conn, item.clone(), digest_type, digest)?; for meta_plugin in meta_plugins.iter_mut() { let meta_name = meta_plugin.meta_name(); // TODO: Add error handling instead of unwrap. let meta_value = meta_plugin.finalize().unwrap(); store_item_meta_value(conn, item.clone(), meta_name, meta_value)?; } db::update_item(conn, item.clone())?; Ok(()) }