queuectl/src/compression.rs

use async_compression::tokio::bufread::{GzipDecoder, XzDecoder, ZstdDecoder};
use clap::ArgEnum;
use tokio::io::{AsyncRead, BufReader};

#[derive(Copy, Clone, PartialEq, Ord, PartialOrd, Eq, ArgEnum, Debug)]
pub enum CompressionMode {
    AUTO,
    NONE,
    ZSTD,
    GZIP,
    XZ,
}

impl CompressionMode {
    pub fn parse(filename: &str) -> Self {
        if filename.to_ascii_lowercase().ends_with("zst") {
            return Self::ZSTD;
        }
        if filename.to_ascii_lowercase().ends_with("zstd") {
            return Self::ZSTD;
        }
        if filename.to_ascii_lowercase().ends_with("xz") {
            return Self::XZ;
        }
        if filename.to_ascii_lowercase().ends_with("gz") {
            return Self::GZIP;
        }

        return Self::NONE;
    }
}

pub async fn get_decompressed_reader(
    compression_mode: CompressionMode,
    reader: Box<dyn AsyncRead + Unpin>,
) -> anyhow::Result<Box<dyn AsyncRead + Unpin>> {
    match compression_mode {
        CompressionMode::AUTO => unreachable!(),
        CompressionMode::NONE => Ok(Box::new(reader)),
        CompressionMode::ZSTD => {
            let reader = BufReader::new(reader);
            Ok(Box::new(ZstdDecoder::new(reader)))
        }
        CompressionMode::GZIP => {
            let reader = BufReader::new(reader);
            Ok(Box::new(GzipDecoder::new(reader)))
        }
        CompressionMode::XZ => {
            let reader = BufReader::new(reader);
            Ok(Box::new(XzDecoder::new(reader)))
        }
    }
}