ml-cellar 0.2.0

CLI of ML model registry for minimum MLOps
Documentation
use serde::{Deserialize, Serialize};
use std::env;
use std::path::{Path, PathBuf};

/// Configuration for the ml-cellar repository.
#[derive(Debug, Deserialize, Serialize, Default)]
#[serde(default)]
pub struct CellarConfig {
    pub ml_cellar: MlCellarConfig,
    pub aws: AWSConfig,
}

/// Configuration specific to ml-cellar behavior.
#[derive(Debug, Deserialize, Serialize, Default)]
#[serde(default)]
pub struct MlCellarConfig {
    /// Whether to use a custom transfer agent for Git LFS (e.g., for AWS S3 backend).
    ///
    /// When enabled, ml-cellar can use alternative storage backends instead of
    /// the default Git LFS storage.
    pub use_custom_transfer_agent: bool,
}

/// Configuration for AWS integration.
#[derive(Debug, Deserialize, Serialize, Default)]
#[serde(default)]
pub struct AWSConfig {
    /// AWS profile name to use for authentication.
    /// If specified, ml-cellar will use credentials from this AWS CLI profile.
    pub profile: Option<String>,
}

/// Loads the cellar configuration by searching for `.mlcellar.toml` in the directory tree.
///
/// This function searches upward from the given path through parent directories
/// until it finds a `.mlcellar.toml` file. This allows ml-cellar commands to be
/// run from any subdirectory within a repository.
///
/// # Arguments
///
/// - `path` - The starting path to search from (can be a file or directory)
///
/// # Returns
///
/// A tuple `(CellarConfig, PathBuf)` where:
/// - `CellarConfig`: The parsed configuration from `.mlcellar.toml`
/// - `PathBuf`: The absolute path to the directory containing `.mlcellar.toml`
///   (this is the root of the ml-cellar repository)
///
/// # Panics
///
/// Panics if:
/// - `.mlcellar.toml` is not found in any parent directory up to the filesystem root
/// - The configuration file cannot be read
/// - The TOML content cannot be parsed
///
pub fn load_cellar_config(path: &Path) -> (CellarConfig, PathBuf) {
    // Set directory to start searching
    let relative_dir = if path.is_dir() {
        path.to_path_buf()
    } else {
        path.parent().unwrap().to_path_buf()
    };

    // Convert to absolute path
    let mut absolute_dir = if relative_dir.is_absolute() {
        relative_dir
    } else {
        env::current_dir().unwrap().join(relative_dir)
    };

    // Search for .mlcellar.toml in the directory and its parents
    loop {
        let candidate = absolute_dir.join(".mlcellar.toml");
        if candidate.is_file() {
            // Found .mlcellar.toml, read and parse it
            log::info!("Loading config from {:?}", candidate);
            let config_content = std::fs::read_to_string(&candidate).unwrap();
            return (toml::from_str(&config_content).unwrap(), absolute_dir);
        }

        match absolute_dir.parent() {
            Some(parent) => absolute_dir = parent.to_path_buf(),
            None => {
                log::error!(
                    ".mlcellar.toml not found; reached filesystem root at {:?}\n\
                     Please ensure that .mlcellar.toml exists in the directory tree starting from the directory of the provided path.",
                    absolute_dir
                );
                panic!(
                    ".mlcellar.toml not found in the directory tree starting from {:?}",
                    path
                );
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::TempDir;
    use toml::to_string_pretty;

    #[test]
    fn test_load_cellar_config() {
        let temp = TempDir::new().unwrap();
        let root_directory = temp.path();

        let toml_str = to_string_pretty(&CellarConfig::default())
            .expect("failed to serialize cellar config to TOML");
        std::fs::write(root_directory.join(".mlcellar.toml"), toml_str)
            .expect("failed to write .mlcellar.toml");

        let (config, config_dir) = load_cellar_config(root_directory);
        assert_eq!(config_dir, root_directory);
        assert!(!config.ml_cellar.use_custom_transfer_agent);
    }
}