Files
rustytorch/crates/production/rtx-config/src/loader.rs
T
2026-03-04 00:08:42 +00:00

807 lines
25 KiB
Rust

//! Configuration loading from various sources.
use crate::{ConfigError, ConfigResult, ConfigValue};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::Path;
use std::time::Duration;
use tracing::{debug, warn};
/// Configuration format types.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ConfigFormat {
Json,
Yaml,
Toml,
Env,
}
/// Configuration loader settings.
#[derive(Debug, Clone)]
pub struct LoaderConfig {
/// Environment variable prefix (e.g., "RTX_" for RTX_DATABASE_URL)
pub env_prefix: Option<String>,
/// Case sensitivity for environment variables
pub env_case_sensitive: bool,
/// Separator for nested keys in environment variables (e.g., "__" for RTX_DB__URL)
pub env_separator: String,
/// Whether to ignore missing optional files
pub ignore_missing_files: bool,
/// Default timeout for remote configuration loading
pub remote_timeout: Duration,
/// Maximum file size for configuration files (in bytes)
pub max_file_size: usize,
}
impl Default for LoaderConfig {
fn default() -> Self {
Self {
env_prefix: Some("RTX_".to_string()),
env_case_sensitive: false,
env_separator: "__".to_string(),
ignore_missing_files: true,
remote_timeout: Duration::from_secs(30),
max_file_size: 10 * 1024 * 1024, // 10MB
}
}
}
/// Configuration loader for different sources.
#[derive(Debug, Clone)]
pub struct ConfigLoader {
config: LoaderConfig,
}
impl ConfigLoader {
/// Create a new configuration loader.
pub fn new(config: LoaderConfig) -> Self {
Self { config }
}
/// Load configuration from environment variables.
pub async fn load_environment(&self) -> ConfigResult<HashMap<String, ConfigValue>> {
let mut values = HashMap::new();
for (key, value) in std::env::vars() {
if let Some(config_key) = self.process_env_var(&key) {
let config_value = self.parse_env_value(&value)?;
values.insert(config_key, config_value);
}
}
debug!(
"Loaded {} configuration values from environment",
values.len()
);
Ok(values)
}
/// Load configuration from a file.
pub async fn load_file(
&self,
path: &Path,
format: ConfigFormat,
) -> ConfigResult<HashMap<String, ConfigValue>> {
if !path.exists() {
if self.config.ignore_missing_files {
debug!("Configuration file not found (ignored): {}", path.display());
return Ok(HashMap::new());
}
return Err(ConfigError::FileNotFound {
path: path.display().to_string(),
});
}
// Check file size
let metadata = tokio::fs::metadata(path).await?;
if metadata.len() > self.config.max_file_size as u64 {
return Err(ConfigError::ValidationFailed {
details: format!(
"Configuration file too large: {} bytes (max: {} bytes)",
metadata.len(),
self.config.max_file_size
),
});
}
let content = tokio::fs::read_to_string(path).await?;
let values =
self.parse_content(&content, format)
.map_err(|e| ConfigError::InvalidFormat {
path: path.display().to_string(),
details: e.to_string(),
})?;
debug!(
"Loaded {} configuration values from file: {}",
values.len(),
path.display()
);
Ok(values)
}
/// Load configuration from a remote URL.
#[cfg(feature = "remote-config")]
pub async fn load_remote(
&self,
url: &str,
headers: HashMap<String, String>,
timeout: Duration,
) -> ConfigResult<HashMap<String, ConfigValue>> {
let client = reqwest::Client::builder()
.timeout(timeout)
.build()
.map_err(|e| ConfigError::remote_error(url, e.to_string()))?;
let mut request = client.get(url);
// Add custom headers
for (key, value) in headers {
request = request.header(&key, &value);
}
let response = request
.send()
.await
.map_err(|e| ConfigError::remote_error(url, e.to_string()))?;
if !response.status().is_success() {
return Err(ConfigError::remote_error(
url,
format!("HTTP error: {}", response.status()),
));
}
let content_type = response
.headers()
.get("content-type")
.and_then(|v| v.to_str().ok())
.unwrap_or("application/json");
let format = self.detect_format_from_content_type(content_type);
let content = response
.text()
.await
.map_err(|e| ConfigError::remote_error(url, e.to_string()))?;
let values = self
.parse_content(&content, format)
.map_err(|e| ConfigError::remote_error(url, format!("Parse error: {}", e)))?;
debug!(
"Loaded {} configuration values from remote URL: {}",
values.len(),
url
);
Ok(values)
}
/// Load configuration from a remote URL.
#[cfg(not(feature = "remote-config"))]
pub async fn load_remote(
&self,
_url: &str,
_headers: HashMap<String, String>,
_timeout: Duration,
) -> ConfigResult<HashMap<String, ConfigValue>> {
Err(ConfigError::remote_error(
_url,
"Remote configuration loading is not enabled. Enable the 'remote-config' feature to use this functionality.".to_string(),
))
}
/// Process environment variable name to configuration key.
fn process_env_var(&self, env_key: &str) -> Option<String> {
let processed_key = if !self.config.env_case_sensitive {
env_key.to_uppercase()
} else {
env_key.to_string()
};
if let Some(ref prefix) = self.config.env_prefix {
let prefix = if !self.config.env_case_sensitive {
prefix.to_uppercase()
} else {
prefix.clone()
};
if processed_key.starts_with(&prefix) {
let config_key = processed_key.strip_prefix(&prefix)?.to_lowercase();
// Convert separator to dots for nested keys
let config_key = config_key.replace(&self.config.env_separator, ".");
Some(config_key)
} else {
None
}
} else {
// No prefix, use the key as-is (converted to lowercase)
Some(
processed_key
.to_lowercase()
.replace(&self.config.env_separator, "."),
)
}
}
/// Parse environment variable value to ConfigValue.
fn parse_env_value(&self, value: &str) -> ConfigResult<ConfigValue> {
// Try to parse as different types
// Boolean
match value.to_lowercase().as_str() {
"true" | "yes" | "1" | "on" => return Ok(ConfigValue::Boolean(true)),
"false" | "no" | "0" | "off" => return Ok(ConfigValue::Boolean(false)),
_ => {}
}
// Integer
if let Ok(int_val) = value.parse::<i64>() {
return Ok(ConfigValue::Integer(int_val));
}
// Float
if let Ok(float_val) = value.parse::<f64>() {
return Ok(ConfigValue::Float(float_val));
}
// JSON (arrays, objects)
if value.starts_with('[') || value.starts_with('{') {
if let Ok(json_val) = serde_json::from_str::<serde_json::Value>(value) {
return Ok(self.json_to_config_value(json_val));
}
}
// Default to string
Ok(ConfigValue::String(value.to_string()))
}
/// Parse configuration content based on format.
fn parse_content(
&self,
content: &str,
format: ConfigFormat,
) -> ConfigResult<HashMap<String, ConfigValue>> {
match format {
ConfigFormat::Json => {
let json_value: serde_json::Value = serde_json::from_str(content)?;
Ok(self.json_to_flat_map(json_value, None))
}
ConfigFormat::Yaml => {
let yaml_value: serde_yaml::Value = serde_yaml::from_str(content)?;
Ok(self.yaml_to_flat_map(yaml_value, None))
}
ConfigFormat::Toml => {
let toml_value: toml::Value = toml::from_str(content)?;
Ok(self.toml_to_flat_map(toml_value, None))
}
ConfigFormat::Env => {
// Parse as .env file format
self.parse_env_file(content)
}
}
}
/// Convert JSON value to ConfigValue.
fn json_to_config_value(&self, value: serde_json::Value) -> ConfigValue {
match value {
serde_json::Value::Null => ConfigValue::Null,
serde_json::Value::Bool(b) => ConfigValue::Boolean(b),
serde_json::Value::Number(n) => {
if let Some(i) = n.as_i64() {
ConfigValue::Integer(i)
} else if let Some(f) = n.as_f64() {
ConfigValue::Float(f)
} else {
ConfigValue::String(n.to_string())
}
}
serde_json::Value::String(s) => ConfigValue::String(s),
serde_json::Value::Array(arr) => {
let config_array = arr
.into_iter()
.map(|v| self.json_to_config_value(v))
.collect();
ConfigValue::Array(config_array)
}
serde_json::Value::Object(obj) => {
let config_object = obj
.into_iter()
.map(|(k, v)| (k, self.json_to_config_value(v)))
.collect();
ConfigValue::Object(config_object)
}
}
}
/// Flatten JSON value to dot-notation keys.
fn json_to_flat_map(
&self,
value: serde_json::Value,
prefix: Option<String>,
) -> HashMap<String, ConfigValue> {
let mut result = HashMap::new();
if let serde_json::Value::Object(obj) = value {
for (key, val) in obj {
let new_key = if let Some(ref p) = prefix {
format!("{}.{}", p, key)
} else {
key
};
match val {
serde_json::Value::Object(_) => {
result.extend(self.json_to_flat_map(val, Some(new_key)));
}
_ => {
result.insert(new_key, self.json_to_config_value(val));
}
}
}
} else {
let key = prefix.unwrap_or_else(|| "root".to_string());
result.insert(key, self.json_to_config_value(value));
}
result
}
/// Convert YAML value to flat map.
fn yaml_to_flat_map(
&self,
value: serde_yaml::Value,
prefix: Option<String>,
) -> HashMap<String, ConfigValue> {
let mut result = HashMap::new();
if let serde_yaml::Value::Mapping(map) = value {
for (key, val) in map {
if let serde_yaml::Value::String(key_str) = key {
let new_key = if let Some(ref p) = prefix {
format!("{}.{}", p, key_str)
} else {
key_str
};
match val {
serde_yaml::Value::Mapping(_) => {
result.extend(self.yaml_to_flat_map(val, Some(new_key)));
}
_ => {
result.insert(new_key, self.yaml_to_config_value(val));
}
}
}
}
} else {
let key = prefix.unwrap_or_else(|| "root".to_string());
result.insert(key, self.yaml_to_config_value(value));
}
result
}
/// Convert YAML value to ConfigValue.
fn yaml_to_config_value(&self, value: serde_yaml::Value) -> ConfigValue {
match value {
serde_yaml::Value::Null => ConfigValue::Null,
serde_yaml::Value::Bool(b) => ConfigValue::Boolean(b),
serde_yaml::Value::Number(n) => {
if let Some(i) = n.as_i64() {
ConfigValue::Integer(i)
} else if let Some(f) = n.as_f64() {
ConfigValue::Float(f)
} else {
ConfigValue::String(n.to_string())
}
}
serde_yaml::Value::String(s) => ConfigValue::String(s),
serde_yaml::Value::Sequence(seq) => {
let config_array = seq
.into_iter()
.map(|v| self.yaml_to_config_value(v))
.collect();
ConfigValue::Array(config_array)
}
serde_yaml::Value::Mapping(map) => {
let mut config_object = HashMap::new();
for (k, v) in map {
if let serde_yaml::Value::String(key_str) = k {
config_object.insert(key_str, self.yaml_to_config_value(v));
}
}
ConfigValue::Object(config_object)
}
_ => {
// For other types, serialize to string using YAML format
match serde_yaml::to_string(&value) {
Ok(s) => ConfigValue::String(s),
Err(_) => ConfigValue::Null,
}
}
}
}
/// Convert TOML value to flat map.
fn toml_to_flat_map(
&self,
value: toml::Value,
prefix: Option<String>,
) -> HashMap<String, ConfigValue> {
let mut result = HashMap::new();
if let toml::Value::Table(table) = value {
for (key, val) in table {
let new_key = if let Some(ref p) = prefix {
format!("{}.{}", p, key)
} else {
key
};
match val {
toml::Value::Table(_) => {
result.extend(self.toml_to_flat_map(val, Some(new_key)));
}
_ => {
result.insert(new_key, self.toml_to_config_value(val));
}
}
}
} else {
let key = prefix.unwrap_or_else(|| "root".to_string());
result.insert(key, self.toml_to_config_value(value));
}
result
}
/// Convert TOML value to ConfigValue.
fn toml_to_config_value(&self, value: toml::Value) -> ConfigValue {
match value {
toml::Value::String(s) => ConfigValue::String(s),
toml::Value::Integer(i) => ConfigValue::Integer(i),
toml::Value::Float(f) => ConfigValue::Float(f),
toml::Value::Boolean(b) => ConfigValue::Boolean(b),
toml::Value::Array(arr) => {
let config_array = arr
.into_iter()
.map(|v| self.toml_to_config_value(v))
.collect();
ConfigValue::Array(config_array)
}
toml::Value::Table(table) => {
let config_object = table
.into_iter()
.map(|(k, v)| (k, self.toml_to_config_value(v)))
.collect();
ConfigValue::Object(config_object)
}
toml::Value::Datetime(dt) => ConfigValue::String(dt.to_string()),
}
}
/// Parse .env file format.
fn parse_env_file(&self, content: &str) -> ConfigResult<HashMap<String, ConfigValue>> {
let mut values = HashMap::new();
for line in content.lines() {
let line = line.trim();
// Skip comments and empty lines
if line.is_empty() || line.starts_with('#') {
continue;
}
// Parse KEY=VALUE format
if let Some((key, value)) = line.split_once('=') {
let key = key.trim();
let value = value.trim();
// Remove quotes if present
let value = if (value.starts_with('"') && value.ends_with('"'))
|| (value.starts_with('\'') && value.ends_with('\''))
{
&value[1..value.len() - 1]
} else {
value
};
if let Some(config_key) = self.process_env_var(key) {
let config_value = self.parse_env_value(value)?;
values.insert(config_key, config_value);
}
} else {
warn!("Ignoring invalid .env line: {}", line);
}
}
Ok(values)
}
/// Detect configuration format from content type.
fn detect_format_from_content_type(&self, content_type: &str) -> ConfigFormat {
if content_type.contains("application/json") {
ConfigFormat::Json
} else if content_type.contains("application/x-yaml") || content_type.contains("text/yaml")
{
ConfigFormat::Yaml
} else if content_type.contains("application/toml") {
ConfigFormat::Toml
} else {
// Default to JSON
ConfigFormat::Json
}
}
}
impl Default for ConfigLoader {
fn default() -> Self {
Self::new(LoaderConfig::default())
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Write;
use tempfile::NamedTempFile;
#[tokio::test]
async fn test_load_json_file() {
let json_content = r#"{
"database": {
"url": "postgres://localhost/test",
"port": 5432
},
"debug": true
}"#;
let mut temp_file = NamedTempFile::new().unwrap();
write!(temp_file, "{}", json_content).unwrap();
let loader = ConfigLoader::default();
let values = loader
.load_file(temp_file.path(), ConfigFormat::Json)
.await
.unwrap();
assert_eq!(
values.get("database.url").unwrap(),
&ConfigValue::String("postgres://localhost/test".to_string())
);
assert_eq!(
values.get("database.port").unwrap(),
&ConfigValue::Integer(5432)
);
assert_eq!(values.get("debug").unwrap(), &ConfigValue::Boolean(true));
}
#[tokio::test]
async fn test_load_yaml_file() {
let yaml_content = r#"
database:
url: postgres://localhost/test
port: 5432
debug: true
features:
- auth
- logging
"#;
let mut temp_file = NamedTempFile::new().unwrap();
write!(temp_file, "{}", yaml_content).unwrap();
let loader = ConfigLoader::default();
let values = loader
.load_file(temp_file.path(), ConfigFormat::Yaml)
.await
.unwrap();
assert_eq!(
values.get("database.url").unwrap(),
&ConfigValue::String("postgres://localhost/test".to_string())
);
assert_eq!(
values.get("database.port").unwrap(),
&ConfigValue::Integer(5432)
);
assert_eq!(values.get("debug").unwrap(), &ConfigValue::Boolean(true));
if let ConfigValue::Array(features) = values.get("features").unwrap() {
assert_eq!(features.len(), 2);
} else {
panic!("Expected features to be an array");
}
}
#[tokio::test]
async fn test_load_toml_file() {
let toml_content = r#"
debug = true
[database]
url = "postgres://localhost/test"
port = 5432
[features]
auth = true
logging = false
"#;
let mut temp_file = NamedTempFile::new().unwrap();
write!(temp_file, "{}", toml_content).unwrap();
let loader = ConfigLoader::default();
let values = loader
.load_file(temp_file.path(), ConfigFormat::Toml)
.await
.unwrap();
assert_eq!(
values.get("database.url").unwrap(),
&ConfigValue::String("postgres://localhost/test".to_string())
);
assert_eq!(
values.get("database.port").unwrap(),
&ConfigValue::Integer(5432)
);
assert_eq!(values.get("debug").unwrap(), &ConfigValue::Boolean(true));
assert_eq!(
values.get("features.auth").unwrap(),
&ConfigValue::Boolean(true)
);
assert_eq!(
values.get("features.logging").unwrap(),
&ConfigValue::Boolean(false)
);
}
#[tokio::test]
async fn test_load_env_file() {
let env_content = r#"
# Database configuration
RTX_DATABASE__URL=postgres://localhost/test
RTX_DATABASE__PORT=5432
RTX_DEBUG=true
# Quoted values
RTX_SECRET_KEY="my-secret-key"
RTX_DESCRIPTION='This is a test'
"#;
let mut temp_file = NamedTempFile::new().unwrap();
write!(temp_file, "{}", env_content).unwrap();
let loader = ConfigLoader::default();
let values = loader
.load_file(temp_file.path(), ConfigFormat::Env)
.await
.unwrap();
assert_eq!(
values.get("database.url").unwrap(),
&ConfigValue::String("postgres://localhost/test".to_string())
);
assert_eq!(
values.get("database.port").unwrap(),
&ConfigValue::Integer(5432)
);
assert_eq!(values.get("debug").unwrap(), &ConfigValue::Boolean(true));
assert_eq!(
values.get("secret_key").unwrap(),
&ConfigValue::String("my-secret-key".to_string())
);
assert_eq!(
values.get("description").unwrap(),
&ConfigValue::String("This is a test".to_string())
);
}
#[test]
fn test_env_var_processing() {
let config = LoaderConfig::default();
let loader = ConfigLoader::new(config);
// Test with prefix
assert_eq!(
loader.process_env_var("RTX_DATABASE_URL"),
Some("database_url".to_string())
);
assert_eq!(
loader.process_env_var("RTX_DB__HOST"),
Some("db.host".to_string())
);
// Test without prefix match
assert_eq!(loader.process_env_var("PATH"), None);
assert_eq!(loader.process_env_var("HOME"), None);
}
#[test]
fn test_env_value_parsing() {
let loader = ConfigLoader::default();
// Boolean values
assert_eq!(
loader.parse_env_value("true").unwrap(),
ConfigValue::Boolean(true)
);
assert_eq!(
loader.parse_env_value("false").unwrap(),
ConfigValue::Boolean(false)
);
assert_eq!(
loader.parse_env_value("yes").unwrap(),
ConfigValue::Boolean(true)
);
assert_eq!(
loader.parse_env_value("no").unwrap(),
ConfigValue::Boolean(false)
);
// Integer values
assert_eq!(
loader.parse_env_value("42").unwrap(),
ConfigValue::Integer(42)
);
assert_eq!(
loader.parse_env_value("-123").unwrap(),
ConfigValue::Integer(-123)
);
// Float values
assert_eq!(
loader.parse_env_value("3.14").unwrap(),
ConfigValue::Float(3.14)
);
// String values
assert_eq!(
loader.parse_env_value("hello").unwrap(),
ConfigValue::String("hello".to_string())
);
// JSON values
let json_result = loader.parse_env_value(r#"["a", "b", "c"]"#).unwrap();
if let ConfigValue::Array(arr) = json_result {
assert_eq!(arr.len(), 3);
} else {
panic!("Expected array");
}
}
#[tokio::test]
async fn test_missing_file_handling() {
let loader = ConfigLoader::new(LoaderConfig {
ignore_missing_files: true,
..Default::default()
});
let result = loader
.load_file(
&std::path::PathBuf::from("nonexistent.json"),
ConfigFormat::Json,
)
.await;
assert!(result.is_ok());
assert!(result.unwrap().is_empty());
let strict_loader = ConfigLoader::new(LoaderConfig {
ignore_missing_files: false,
..Default::default()
});
let result = strict_loader
.load_file(
&std::path::PathBuf::from("nonexistent.json"),
ConfigFormat::Json,
)
.await;
assert!(result.is_err());
assert!(matches!(
result.unwrap_err(),
ConfigError::FileNotFound { .. }
));
}
}