tracel_xtask/

environment.rs

use std::{
    collections::HashMap,
    fmt::{self, Display, Write as _},
    marker::PhantomData,
    path::PathBuf,
};

use strum::{EnumIter, EnumString, IntoEnumIterator as _};

use crate::{group_error, group_info, utils::git};

/// Implicit index which means that index '1' is omitted in display.
#[derive(Clone, Debug, PartialEq, Default)]
pub struct ImplicitIndex;

/// Explicit index which means that index is always in display.
#[derive(Clone, Debug, PartialEq, Default)]
pub struct ExplicitIndex;

/// Style for how to format `{base}{index}`.
pub trait IndexStyle {
    fn format(base: &str, index: u8) -> String;
}

impl IndexStyle for ImplicitIndex {
    fn format(base: &str, index: u8) -> String {
        if index == 1 {
            base.to_string()
        } else {
            format!("{base}{index}")
        }
    }
}

impl IndexStyle for ExplicitIndex {
    fn format(base: &str, index: u8) -> String {
        format!("{base}{index}")
    }
}

#[derive(Clone, Debug, Default, PartialEq)]
pub struct Environment<M = ImplicitIndex> {
    pub name: EnvironmentName,
    pub index: EnvironmentIndex,
    _marker: PhantomData<M>,
}

impl<M> Environment<M> {
    pub fn new(name: EnvironmentName, index: u8) -> Self {
        Self {
            name,
            index: index.into(),
            _marker: PhantomData,
        }
    }

    pub fn index(&self) -> u8 {
        self.index.index
    }
}

impl Environment<ImplicitIndex> {
    /// Turn an non explicit environment into an explicit one.
    /// An explicit environment will always append the index number to its display names.
    /// Whereas a non-explicit one (default) only append the index if it is different than 1.
    pub fn into_explicit(self) -> Environment<ExplicitIndex> {
        Environment {
            name: self.name.clone(),
            index: self.index().into(),
            _marker: PhantomData,
        }
    }
}

impl Environment<ExplicitIndex> {
    pub fn into_implicit(self) -> Environment<ImplicitIndex> {
        Environment {
            name: self.name.clone(),
            index: self.index().into(),
            _marker: PhantomData,
        }
    }
}

impl<M: IndexStyle> Display for Environment<M> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.medium())
    }
}

impl<M: IndexStyle> Environment<M> {
    pub fn long(&self) -> String {
        M::format(self.name.long(), self.index())
    }

    pub fn medium(&self) -> String {
        M::format(self.name.medium(), self.index())
    }

    pub fn short(&self) -> String {
        M::format(&self.name.short().to_string(), self.index())
    }

    /// Return the two .env files for a given family:
    /// - Base: `.env`, `.env.<env_medium>`
    /// - Secrets: `.env.secrets`, `.env.<env_medium>.secrets`
    /// - Infra: `.env.infra`, `.env.<env_medium>.infra`
    /// - InfraSecrets: `.env.infra.secrets`, `.env.<env_medium>.infra.secrets`
    fn dotenv_files_for_family(&self, family: DotEnvFamily) -> [String; 2] {
        let suffix = family.to_string();
        let env_medium = self.medium();
        if suffix.is_empty() {
            // Base
            [".env".to_owned(), format!(".env.{env_medium}")]
        } else {
            // Other families
            [
                format!(".env{suffix}"),
                format!(".env.{env_medium}{suffix}"),
            ]
        }
    }

    /// Backward-compatible helper for env-specific base filename.
    pub fn get_dotenv_filename(&self) -> String {
        // second element of the Base family
        self.dotenv_files_for_family(DotEnvFamily::Base)[1].clone()
    }

    /// Backward-compatible helper for env-specific secrets filename.
    pub fn get_dotenv_secrets_filename(&self) -> String {
        // second element of the Secrets family
        self.dotenv_files_for_family(DotEnvFamily::Secrets)[1].clone()
    }

    /// All possible .env files for this environment, by family.
    /// Order matters: later files override earlier ones.
    pub fn get_env_files(&self) -> Vec<String> {
        DotEnvFamily::iter()
            .flat_map(|family| self.dotenv_files_for_family(family))
            .collect()
    }

    /// Load the .env environment files family.
    pub fn load(&self, prefix: Option<&str>) -> anyhow::Result<()> {
        let files = self.get_env_files();
        for file in files {
            let path = if let Some(p) = prefix {
                PathBuf::from(p).join(&file)
            } else {
                PathBuf::from(&file)
            };
            if path.exists() {
                match dotenvy::from_path(&path) {
                    Ok(_) => {
                        group_info!("loading '{}' file...", path.display());
                    }
                    Err(e) => {
                        group_error!("error while loading '{}' file ({})", path.display(), e);
                    }
                }
            }
        }

        Ok(())
    }

    /// Merge all the .env files of the environment with all variable expanded
    pub fn merge_env_files(&self) -> anyhow::Result<PathBuf> {
        let repo_root = git::git_repo_root_or_cwd()?;
        let files = self.get_env_files();
        // merged set of env vars, the later files override earlier ones
        // we sort keys to have a more deterministic merged file result
        let mut merged: HashMap<String, String> = HashMap::new();
        for filename in files {
            let path = repo_root.join(&filename);
            if !path.exists() {
                eprintln!(
                    "⚠️ Warning: environment file '{}' ({}) not found, skipping...",
                    filename,
                    path.display()
                );
                continue;
            }
            for item in dotenvy::from_path_iter(&path)? {
                let (key, value) = item?;
                std::env::set_var(&key, &value);
                merged.insert(key, value);
            }
        }
        let mut keys: Vec<_> = merged.keys().cloned().collect();
        keys.sort();
        // write merged file
        let mut out = String::new();
        for key in keys {
            let val = &merged[&key];
            writeln!(&mut out, "{key}={val}")?;
        }
        let tmp_path = std::env::temp_dir().join(format!("merged-env-{}.tmp", std::process::id()));
        std::fs::write(&tmp_path, out)?;
        Ok(tmp_path)
    }
}

#[derive(EnumString, EnumIter, Default, Clone, Debug, PartialEq, clap::ValueEnum)]
#[strum(serialize_all = "lowercase")]
pub enum EnvironmentName {
    /// Development environment (alias: dev).
    #[default]
    #[clap(alias = "dev")]
    Development,
    /// Staging environment (alias: stag).
    #[clap(alias = "stag")]
    Staging,
    /// Testing environment (alias: test).
    #[clap(alias = "test")]
    Test,
    /// Production environment (alias: prod).
    #[clap(alias = "prod")]
    Production,
}

impl Display for EnvironmentName {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.medium())
    }
}

impl EnvironmentName {
    pub fn long(&self) -> &'static str {
        match self {
            EnvironmentName::Development => "development",
            EnvironmentName::Staging => "staging",
            EnvironmentName::Test => "test",
            EnvironmentName::Production => "production",
        }
    }

    pub fn medium(&self) -> &'static str {
        match self {
            EnvironmentName::Development => "dev",
            EnvironmentName::Staging => "stag",
            EnvironmentName::Test => "test",
            EnvironmentName::Production => "prod",
        }
    }

    pub fn short(&self) -> char {
        match self {
            EnvironmentName::Development => 'd',
            EnvironmentName::Staging => 's',
            EnvironmentName::Test => 't',
            EnvironmentName::Production => 'p',
        }
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct EnvironmentIndex {
    pub index: u8,
}

impl Default for EnvironmentIndex {
    fn default() -> Self {
        Self { index: 1 }
    }
}

impl Display for EnvironmentIndex {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.index)
    }
}

impl From<u8> for EnvironmentIndex {
    fn from(index: u8) -> Self {
        Self { index }
    }
}

#[derive(EnumString, EnumIter, Clone, Debug, PartialEq, clap::ValueEnum)]
enum DotEnvFamily {
    Base,
    Secrets,
    Infra,
    InfraSecrets,
}

impl Display for DotEnvFamily {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            DotEnvFamily::Base => write!(f, ""),
            DotEnvFamily::Secrets => write!(f, ".secrets"),
            DotEnvFamily::Infra => write!(f, ".infra"),
            DotEnvFamily::InfraSecrets => write!(f, ".infra.secrets"),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use rstest::rstest;
    use serial_test::serial;
    use std::env;

    // For tests we always use the implicit style
    type TestEnv = Environment<ImplicitIndex>;

    fn expected_vars(env: &TestEnv) -> Vec<(String, String)> {
        let suffix = match env.name {
            EnvironmentName::Development => "DEV",
            EnvironmentName::Staging => "STAG",
            EnvironmentName::Test => "TEST",
            EnvironmentName::Production => "PROD",
        };

        vec![
            ("FROM_DOTENV".to_string(), ".env".to_string()),
            (
                format!("FROM_DOTENV_{suffix}").to_string(),
                env.get_dotenv_filename(),
            ),
            (
                format!("FROM_DOTENV_{suffix}_SECRETS").to_string(),
                env.get_dotenv_secrets_filename(),
            ),
        ]
    }

    #[rstest]
    #[case::dev(TestEnv::new(EnvironmentName::Development, 1))]
    #[case::stag(TestEnv::new(EnvironmentName::Staging, 1))]
    #[case::test(TestEnv::new(EnvironmentName::Test, 1))]
    #[case::prod(TestEnv::new(EnvironmentName::Production, 1))]
    #[serial]
    fn test_environment_load(#[case] env: TestEnv) {
        // Remove possible prior values
        for (key, _) in expected_vars(&env) {
            env::remove_var(key);
        }

        // Run the actual function under test
        env.load(Some("../.."))
            .expect("Environment load should succeed");

        // Assert each expected env var is present and has the correct value
        for (key, expected_value) in expected_vars(&env) {
            let actual_value =
                env::var(&key).unwrap_or_else(|_| panic!("Missing expected env var: {key}"));
            assert_eq!(
                actual_value, expected_value,
                "Environment variable {key} should be set to {expected_value} but was {actual_value}"
            );
        }
    }

    #[rstest]
    #[case::dev(TestEnv::new(EnvironmentName::Development, 1))]
    #[case::stag(TestEnv::new(EnvironmentName::Staging, 1))]
    #[case::test(TestEnv::new(EnvironmentName::Test, 1))]
    #[case::prod(TestEnv::new(EnvironmentName::Production, 1))]
    #[serial]
    fn test_environment_merge_env_files(#[case] env: TestEnv) {
        // Make sure we start from a clean state
        for (key, _) in expected_vars(&env) {
            env::remove_var(key);
        }
        // Generate the merged env file
        let merged_path = env
            .merge_env_files()
            .expect("merge_env_files should succeed");
        assert!(
            merged_path.exists(),
            "Merged env file should exist at {}",
            merged_path.display()
        );
        // Parse the merged file as a .env file again
        let mut merged_map: std::collections::HashMap<String, String> =
            std::collections::HashMap::new();
        for item in
            dotenvy::from_path_iter(&merged_path).expect("Reading merged env file should succeed")
        {
            let (key, value) = item.expect("Parsing key/value from merged env file should succeed");
            merged_map.insert(key, value);
        }
        // All the vars we expect from the individual files must be present
        for (key, expected_value) in expected_vars(&env) {
            let actual_value = merged_map
                .get(&key)
                .unwrap_or_else(|| panic!("Missing expected merged env var: {key}"));
            assert_eq!(
                actual_value, &expected_value,
                "Merged env var {key} should be {expected_value} but was {actual_value}"
            );
        }
    }

    #[test]
    #[serial]
    fn test_environment_merge_env_files_expansion() {
        let env = Environment::<ImplicitIndex>::new(EnvironmentName::Staging, 1);
        // Clean any prior values that could interfere
        env::remove_var("LOG_LEVEL_TEST");
        env::remove_var("RUST_LOG_TEST");
        env::remove_var("RUST_LOG_STAG_TEST");

        let merged_path = env
            .merge_env_files()
            .expect("merge_env_files should succeed");
        let mut merged_map: std::collections::HashMap<String, String> =
            std::collections::HashMap::new();
        for item in
            dotenvy::from_path_iter(&merged_path).expect("Reading merged env file should succeed")
        {
            let (key, value) = item.expect("Parsing key/value from merged env file should succeed");
            merged_map.insert(key, value);
        }

        let log_level = merged_map
            .get("LOG_LEVEL_TEST")
            .expect("LOG_LEVEL_TEST should be present in merged env file");
        let rust_log = merged_map
            .get("RUST_LOG_TEST")
            .expect("RUST_LOG_TEST should be present in merged env file");

        // 1) We should not see the raw placeholder anymore
        assert!(
            !rust_log.contains("${LOG_LEVEL_TEST}"),
            "RUST_LOG_TEST should not contain the raw placeholder '${{LOG_LEVEL}}', got: {rust_log}"
        );
        // 2) The expanded LOG_LEVEL_TEST value should appear in RUST_LOG_TEST
        assert!(
            rust_log.contains(log_level),
            "RUST_LOG_TEST should contain the expanded LOG_LEVEL_TEST value; LOG_LEVEL_TEST={log_level}, RUST_LOG_TEST={rust_log}"
        );
        // Cross-file expansion with RUST_LOG_STAG_TEST that references LOG_LEVEL_TEST from base .env
        let rust_log_stag = merged_map
            .get("RUST_LOG_STAG_TEST")
            .expect("RUST_LOG_STAG_TEST should be present in merged env file");
        // 3) No raw placeholder in the cross-file value either
        assert!(
            !rust_log_stag.contains("${LOG_LEVEL_TEST}"),
            "RUST_LOG_STAG_TEST should not contain the raw placeholder '${{LOG_LEVEL_TEST}}', got: {rust_log_stag}"
        );
        // 4) The expanded LOG_LEVEL_TEST value should appear in RUST_LOG_STAG_TEST
        assert!(
            rust_log_stag.contains(log_level),
            "RUST_LOG_STAG_TEST should contain the expanded LOG_LEVEL_TEST value; LOG_LEVEL_TEST={log_level}, RUST_LOG_STAG_TEST={rust_log_stag}"
        );
    }
}