queuectl/src/connector.rs

use anyhow::{anyhow, Context};
use async_read_progress::*;
use futures::stream::StreamExt;
use tokio::fs::File;
use tokio::io::AsyncRead;
use tokio_util::io::StreamReader;
use tracing::info;

use crate::args::ARGS;
use crate::compression;
use crate::pbar;

pub async fn open_http(
    url: &str,
) -> anyhow::Result<(Box<dyn AsyncRead + Unpin>, Option<u64>, String)> {
    info!("Requesting url {:?}...", url);
    let client = reqwest::Client::new();
    let resp = client
        .get(url)
        .send()
        .await
        .context("http client get request")?;
    let filename = String::from(resp.url().path_segments().unwrap().last().unwrap());
    info!("Completed request with code {}!", resp.status());
    if !resp.status().is_success() {
        return Err(anyhow!("incorrect status code attempting to retrieve file"));
    }
    let size = resp.content_length();
    let stream = resp.bytes_stream();
    let stream = stream.map(|v| v.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e)));
    Ok((Box::new(StreamReader::new(stream)), size, filename))
}

pub async fn open(
    update_progressbar: bool,
) -> anyhow::Result<(Box<dyn AsyncRead + Unpin>, Option<u64>)> {
    let (reader, total_bytes, filename): (Box<dyn AsyncRead + Unpin>, Option<u64>, String) =
        if ARGS.input == "-" {
            info!("Reading from stdin...");
            (Box::new(tokio::io::stdin()), None, "dummy.txt".to_string())
        } else {
            if ARGS.input.starts_with("http") {
                open_http(&ARGS.input).await?
            } else {
                info!("Opening file {}...", ARGS.input);
                let file = File::open(&ARGS.input).await.context("open input file")?;
                let meta = file.metadata().await.context("read input file metadata")?;
                (Box::new(file), Some(meta.len()), ARGS.input.clone())
            }
        };

    //Hook into the progress bar.
    let reader = {
        Box::new(
            reader.report_progress(std::time::Duration::from_millis(20), move |bytes_read| {
                if update_progressbar {
                    pbar::update_progress(bytes_read as u64);
                }
            }),
        )
    };

    let compression_mode = if ARGS.compression == compression::CompressionMode::AUTO {
        info!("Attempting to guess compression mode...");
        compression::CompressionMode::parse(&filename)
    } else {
        ARGS.compression
    };
    info!("Using compression mode {:?}.", compression_mode);

    info!("Attempting decompression...");
    let reader = compression::get_decompressed_reader(compression_mode, reader).await?;
    info!("Ok!");
    Ok((reader, total_bytes))
}