use crate::storage::error::StorageError;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum StorageType {
Local,
Aws,
Azure,
Gcs,
Hdfs,
Http,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StorageConfig {
#[serde(rename = "type")]
pub storage_type: StorageType,
#[serde(default)]
pub options: HashMap<String, String>,
}
impl StorageConfig {
pub fn new(storage_type: impl Into<String>) -> Self {
let storage_type_str: String = storage_type.into();
Self::try_new(&storage_type_str)
.unwrap_or_else(|_| panic!("Unknown storage type: {}", storage_type_str))
}
pub fn try_new(storage_type: impl Into<String>) -> Result<Self, StorageError> {
let storage_type_str = storage_type.into();
let storage_type = match storage_type_str.to_lowercase().as_str() {
"local" => StorageType::Local,
"aws" | "s3" => StorageType::Aws,
"azure" => StorageType::Azure,
"gcs" | "gcp" => StorageType::Gcs,
"hdfs" => StorageType::Hdfs,
"http" | "webdav" => StorageType::Http,
_ => {
return Err(StorageError::ConfigError(format!(
"Unknown storage type: {}",
storage_type_str
)))
}
};
Ok(Self {
storage_type,
options: Self::default_options(),
})
}
pub fn local() -> Self {
Self {
storage_type: StorageType::Local,
options: Self::default_options(),
}
}
pub fn aws() -> Self {
Self {
storage_type: StorageType::Aws,
options: Self::default_options(),
}
}
pub fn azure() -> Self {
Self {
storage_type: StorageType::Azure,
options: Self::default_options(),
}
}
pub fn gcs() -> Self {
Self {
storage_type: StorageType::Gcs,
options: HashMap::new(),
}
}
pub fn hdfs() -> Self {
Self {
storage_type: StorageType::Hdfs,
options: Self::default_options(),
}
}
pub fn http() -> Self {
Self {
storage_type: StorageType::Http,
options: Self::default_options(),
}
}
pub fn default_options() -> HashMap<String, String> {
[
("timeout", "1200"),
("connect_timeout", "30"),
("max_retries", "20"),
("retry_timeout", "1200"),
("pool_idle_timeout", "15"),
("pool_max_idle_per_host", "5"),
]
.iter()
.map(|(k, v)| (k.to_string(), v.to_string()))
.collect()
}
pub fn with_option(
mut self,
key: impl Into<String> + Clone,
value: impl Into<String> + Clone,
) -> Self {
self.options.insert(key.into(), value.into());
self
}
pub fn with_options(mut self, options: HashMap<String, String>) -> Self {
self.options.extend(options);
self
}
pub fn get_option(&self, key: &str) -> Option<&String> {
self.options.get(key)
}
pub fn storage_type_str(&self) -> &str {
match self.storage_type {
StorageType::Local => "local",
StorageType::Aws => "aws",
StorageType::Azure => "azure",
StorageType::Gcs => "gcs",
StorageType::Hdfs => "hdfs",
StorageType::Http => "http",
}
}
}
impl From<StorageConfig> for String {
fn from(config: StorageConfig) -> Self {
config.storage_type_str().to_string()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_storage_type_serialization() {
let local = StorageType::Local;
let aws = StorageType::Aws;
let azure = StorageType::Azure;
let gcs = StorageType::Gcs;
let hdfs = StorageType::Hdfs;
let http = StorageType::Http;
assert_eq!(serde_json::to_string(&local).unwrap(), "\"local\"");
assert_eq!(serde_json::to_string(&aws).unwrap(), "\"aws\"");
assert_eq!(serde_json::to_string(&azure).unwrap(), "\"azure\"");
assert_eq!(serde_json::to_string(&gcs).unwrap(), "\"gcs\"");
assert_eq!(serde_json::to_string(&hdfs).unwrap(), "\"hdfs\"");
assert_eq!(serde_json::to_string(&http).unwrap(), "\"http\"");
}
#[test]
fn test_storage_type_deserialization() {
let local: StorageType = serde_json::from_str("\"local\"").unwrap();
let aws: StorageType = serde_json::from_str("\"aws\"").unwrap();
let azure: StorageType = serde_json::from_str("\"azure\"").unwrap();
let gcs: StorageType = serde_json::from_str("\"gcs\"").unwrap();
let hdfs: StorageType = serde_json::from_str("\"hdfs\"").unwrap();
let http: StorageType = serde_json::from_str("\"http\"").unwrap();
assert_eq!(local, StorageType::Local);
assert_eq!(aws, StorageType::Aws);
assert_eq!(azure, StorageType::Azure);
assert_eq!(gcs, StorageType::Gcs);
assert_eq!(hdfs, StorageType::Hdfs);
assert_eq!(http, StorageType::Http);
}
#[test]
fn test_storage_config_new_local() {
let config = StorageConfig::new("local");
assert_eq!(config.storage_type, StorageType::Local);
assert!(!config.options.is_empty());
assert_eq!(config.storage_type_str(), "local");
}
#[test]
fn test_storage_config_new_aws() {
let config1 = StorageConfig::new("aws");
let config2 = StorageConfig::new("s3");
let config3 = StorageConfig::new("AWS");
assert_eq!(config1.storage_type, StorageType::Aws);
assert_eq!(config2.storage_type, StorageType::Aws);
assert_eq!(config3.storage_type, StorageType::Aws);
assert_eq!(config1.storage_type_str(), "aws");
}
#[test]
fn test_storage_config_new_azure() {
let config = StorageConfig::new("azure");
assert_eq!(config.storage_type, StorageType::Azure);
assert_eq!(config.storage_type_str(), "azure");
}
#[test]
fn test_storage_config_new_gcs() {
let config1 = StorageConfig::new("gcs");
let config2 = StorageConfig::new("gcp");
assert_eq!(config1.storage_type, StorageType::Gcs);
assert_eq!(config2.storage_type, StorageType::Gcs);
assert_eq!(config1.storage_type_str(), "gcs");
}
#[test]
#[should_panic(expected = "Unknown storage type")]
fn test_storage_config_new_invalid() {
StorageConfig::new("invalid");
}
#[test]
fn test_storage_config_try_new_valid() {
let config = StorageConfig::try_new("aws").unwrap();
assert_eq!(config.storage_type, StorageType::Aws);
let config = StorageConfig::try_new("s3").unwrap();
assert_eq!(config.storage_type, StorageType::Aws);
let config = StorageConfig::try_new("local").unwrap();
assert_eq!(config.storage_type, StorageType::Local);
let config = StorageConfig::try_new("azure").unwrap();
assert_eq!(config.storage_type, StorageType::Azure);
let config = StorageConfig::try_new("gcs").unwrap();
assert_eq!(config.storage_type, StorageType::Gcs);
let config = StorageConfig::try_new("gcp").unwrap();
assert_eq!(config.storage_type, StorageType::Gcs);
let config = StorageConfig::try_new("hdfs").unwrap();
assert_eq!(config.storage_type, StorageType::Hdfs);
let config = StorageConfig::try_new("http").unwrap();
assert_eq!(config.storage_type, StorageType::Http);
let config = StorageConfig::try_new("webdav").unwrap();
assert_eq!(config.storage_type, StorageType::Http);
}
#[test]
fn test_storage_config_try_new_invalid() {
let result = StorageConfig::try_new("invalid");
assert!(result.is_err());
let err = result.unwrap_err();
assert!(err.to_string().contains("Unknown storage type"));
assert!(err.to_string().contains("invalid"));
}
#[test]
fn test_storage_config_local() {
let config = StorageConfig::local();
assert_eq!(config.storage_type, StorageType::Local);
assert!(!config.options.is_empty());
}
#[test]
fn test_storage_config_aws() {
let config = StorageConfig::aws();
assert_eq!(config.storage_type, StorageType::Aws);
assert!(!config.options.is_empty());
}
#[test]
fn test_storage_config_azure() {
let config = StorageConfig::azure();
assert_eq!(config.storage_type, StorageType::Azure);
assert!(!config.options.is_empty());
}
#[test]
fn test_storage_config_gcs() {
let config = StorageConfig::gcs();
assert_eq!(config.storage_type, StorageType::Gcs);
assert!(config.options.is_empty());
}
#[test]
fn test_storage_config_hdfs() {
let config = StorageConfig::hdfs();
assert_eq!(config.storage_type, StorageType::Hdfs);
assert!(!config.options.is_empty());
assert_eq!(config.storage_type_str(), "hdfs");
}
#[test]
fn test_storage_config_http() {
let config = StorageConfig::http();
assert_eq!(config.storage_type, StorageType::Http);
assert!(!config.options.is_empty());
assert_eq!(config.storage_type_str(), "http");
}
#[test]
fn test_default_options() {
let options = StorageConfig::default_options();
assert_eq!(options.get("timeout"), Some(&"1200".to_string()));
assert_eq!(options.get("connect_timeout"), Some(&"30".to_string()));
assert_eq!(options.get("max_retries"), Some(&"20".to_string()));
assert_eq!(options.get("retry_timeout"), Some(&"1200".to_string()));
assert_eq!(options.get("pool_idle_timeout"), Some(&"15".to_string()));
assert_eq!(
options.get("pool_max_idle_per_host"),
Some(&"5".to_string())
);
}
#[test]
fn test_with_option() {
let config = StorageConfig::local()
.with_option("path", "/tmp/data")
.with_option("custom_key", "custom_value");
assert_eq!(config.get_option("path"), Some(&"/tmp/data".to_string()));
assert_eq!(
config.get_option("custom_key"),
Some(&"custom_value".to_string())
);
}
#[test]
fn test_with_options() {
let mut custom_options = HashMap::new();
custom_options.insert("bucket".to_string(), "my-bucket".to_string());
custom_options.insert("region".to_string(), "us-east-1".to_string());
let config = StorageConfig::aws().with_options(custom_options);
assert_eq!(config.get_option("bucket"), Some(&"my-bucket".to_string()));
assert_eq!(config.get_option("region"), Some(&"us-east-1".to_string()));
assert_eq!(config.get_option("timeout"), Some(&"1200".to_string()));
}
#[test]
fn test_get_option() {
let config = StorageConfig::local().with_option("path", "/tmp/data");
assert_eq!(config.get_option("path"), Some(&"/tmp/data".to_string()));
assert_eq!(config.get_option("nonexistent"), None);
}
#[test]
fn test_storage_type_str() {
assert_eq!(StorageConfig::local().storage_type_str(), "local");
assert_eq!(StorageConfig::aws().storage_type_str(), "aws");
assert_eq!(StorageConfig::azure().storage_type_str(), "azure");
assert_eq!(StorageConfig::gcs().storage_type_str(), "gcs");
assert_eq!(StorageConfig::hdfs().storage_type_str(), "hdfs");
assert_eq!(StorageConfig::http().storage_type_str(), "http");
}
#[test]
fn test_from_storage_config_to_string() {
let local_str: String = StorageConfig::local().into();
let aws_str: String = StorageConfig::aws().into();
let azure_str: String = StorageConfig::azure().into();
let gcs_str: String = StorageConfig::gcs().into();
let hdfs_str: String = StorageConfig::hdfs().into();
let http_str: String = StorageConfig::http().into();
assert_eq!(local_str, "local");
assert_eq!(aws_str, "aws");
assert_eq!(azure_str, "azure");
assert_eq!(gcs_str, "gcs");
assert_eq!(hdfs_str, "hdfs");
assert_eq!(http_str, "http");
}
#[test]
fn test_method_chaining() {
let config = StorageConfig::aws()
.with_option("bucket", "my-bucket")
.with_option("region", "us-west-2")
.with_option("access_key_id", "AxxxxxxxxxNN7EXAMPLE");
assert_eq!(config.storage_type, StorageType::Aws);
assert_eq!(config.get_option("bucket"), Some(&"my-bucket".to_string()));
assert_eq!(config.get_option("region"), Some(&"us-west-2".to_string()));
assert_eq!(
config.get_option("access_key_id"),
Some(&"AxxxxxxxxxNN7EXAMPLE".to_string())
);
}
#[test]
fn test_config_serialization() {
let config = StorageConfig::aws()
.with_option("bucket", "test-bucket")
.with_option("region", "us-east-1");
let json = serde_json::to_string(&config).unwrap();
assert!(json.contains("\"type\":\"aws\""));
assert!(json.contains("\"bucket\""));
assert!(json.contains("\"region\""));
}
#[test]
fn test_config_deserialization() {
let json = r#"{"type":"aws","options":{"bucket":"test-bucket","region":"us-east-1"}}"#;
let config: StorageConfig = serde_json::from_str(json).unwrap();
assert_eq!(config.storage_type, StorageType::Aws);
assert_eq!(
config.get_option("bucket"),
Some(&"test-bucket".to_string())
);
assert_eq!(config.get_option("region"), Some(&"us-east-1".to_string()));
}
#[test]
fn test_option_override() {
let config = StorageConfig::local()
.with_option("timeout", "600")
.with_option("timeout", "900");
assert_eq!(config.get_option("timeout"), Some(&"900".to_string()));
}
#[test]
fn test_clone() {
let config1 = StorageConfig::aws().with_option("bucket", "my-bucket");
let config2 = config1.clone();
assert_eq!(config1.storage_type, config2.storage_type);
assert_eq!(config1.get_option("bucket"), config2.get_option("bucket"));
}
}