keep/src/modes/save.rs

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<i64>,
    tags: &mut Vec<String>,
    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<MetaPluginType> = 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<Box<dyn MetaPlugin>> = 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<String, String> = 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<dyn Write> =
        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(())
}