cargo_lambda_metadata/
env.rs

1use clap::{ArgAction, Args, ValueHint};
2use env_file_reader::read_file;
3use miette::Result;
4use serde::{Deserialize, Serialize, ser::SerializeStruct};
5use std::{collections::HashMap, path::PathBuf};
6
7use crate::{cargo::deserialize_vec_or_map, error::MetadataError};
8
9pub type Environment = HashMap<String, String>;
10
11#[derive(Args, Clone, Debug, Default, Deserialize, Serialize)]
12pub struct EnvOptions {
13    /// Option to add one or many environment variables,
14    /// allows multiple repetitions (--env-var KEY=VALUE --env-var OTHER=NEW-VALUE).
15    /// It also allows to set a list of environment variables separated by commas
16    /// (e.g. --env-var KEY=VALUE,OTHER=NEW-VALUE).
17    #[arg(long, value_delimiter = ',', action = ArgAction::Append, visible_alias = "env-vars")]
18    #[serde(default, alias = "env", deserialize_with = "deserialize_vec_or_map")]
19    pub env_var: Option<Vec<String>>,
20
21    /// Read environment variables from a file.
22    /// Variables are separated by new lines in KEY=VALUE format.
23    #[arg(long, value_hint = ValueHint::FilePath)]
24    #[serde(default)]
25    pub env_file: Option<PathBuf>,
26}
27
28impl EnvOptions {
29    pub fn lambda_environment(
30        &self,
31        base: &HashMap<String, String>,
32    ) -> Result<Environment, MetadataError> {
33        lambda_environment(Some(base), &self.env_file, self.env_var.as_ref())
34    }
35
36    pub fn count_fields(&self) -> usize {
37        self.env_var.is_some() as usize + self.env_file.is_some() as usize
38    }
39
40    pub fn serialize_fields<S>(
41        &self,
42        state: &mut <S as serde::Serializer>::SerializeStruct,
43    ) -> Result<(), S::Error>
44    where
45        S: serde::Serializer,
46    {
47        if let Some(env_var) = &self.env_var {
48            state.serialize_field("env_var", env_var)?;
49        }
50        if let Some(env_file) = &self.env_file {
51            state.serialize_field("env_file", env_file)?;
52        }
53        Ok(())
54    }
55}
56
57pub(crate) fn lambda_environment(
58    base: Option<&HashMap<String, String>>,
59    env_file: &Option<PathBuf>,
60    vars: Option<&Vec<String>>,
61) -> Result<Environment, MetadataError> {
62    let mut env = HashMap::new();
63
64    if let Some(base) = base.cloned() {
65        env.extend(base);
66    }
67
68    if let Some(path) = env_file {
69        if path.is_file() {
70            let env_variables =
71                read_file(path).map_err(|e| MetadataError::InvalidEnvFile(path.into(), e))?;
72            for (key, value) in env_variables {
73                env.insert(key, value);
74            }
75        }
76    }
77
78    if let Some(vars) = vars {
79        for var in vars {
80            let (key, value) = extract_var(var)?;
81            env.insert(key.to_string(), value.to_string());
82        }
83    }
84
85    Ok(env)
86}
87
88fn extract_var(line: &str) -> Result<(&str, &str), MetadataError> {
89    let mut iter = line.trim().splitn(2, '=');
90
91    let key = iter
92        .next()
93        .map(|s| s.trim())
94        .ok_or_else(|| MetadataError::InvalidEnvVar(line.into()))?;
95    if key.is_empty() {
96        Err(MetadataError::InvalidEnvVar(line.into()))?;
97    }
98
99    let value = iter
100        .next()
101        .map(|s| s.trim())
102        .ok_or_else(|| MetadataError::InvalidEnvVar(line.into()))?;
103    if value.is_empty() {
104        Err(MetadataError::InvalidEnvVar(line.into()))?;
105    }
106
107    Ok((key, value))
108}
109
110#[cfg(test)]
111mod test {
112    use std::env::temp_dir;
113
114    use super::*;
115
116    #[test]
117    fn test_extract_var() {
118        let (k, v) = extract_var("FOO=BAR").unwrap();
119        assert_eq!("FOO", k);
120        assert_eq!("BAR", v);
121
122        let (k, v) = extract_var(" FOO = BAR ").unwrap();
123        assert_eq!("FOO", k);
124        assert_eq!("BAR", v);
125
126        extract_var("=BAR").expect_err("missing key");
127        extract_var("FOO=").expect_err("missing value");
128        extract_var("  ").expect_err("missing variable");
129    }
130
131    #[test]
132    fn test_empty_environment() {
133        let env = lambda_environment(None, &None, None).unwrap();
134        assert!(env.is_empty());
135    }
136
137    #[test]
138    fn test_base_environment() {
139        let mut base = HashMap::new();
140        base.insert("FOO".into(), "BAR".into());
141        let env = lambda_environment(Some(&base), &None, None).unwrap();
142
143        assert_eq!("BAR".to_string(), env["FOO"]);
144    }
145
146    #[test]
147    fn test_environment_with_flags() {
148        let mut base = HashMap::new();
149        base.insert("FOO".into(), "BAR".into());
150
151        let flags = vec!["FOO=QUX".to_string(), "BAZ=QUUX".to_string()];
152        let env = lambda_environment(Some(&base), &None, Some(&flags)).unwrap();
153
154        assert_eq!("QUX".to_string(), env["FOO"]);
155        assert_eq!("QUUX".to_string(), env["BAZ"]);
156    }
157
158    #[test]
159    fn test_environment_with_file() {
160        let file = temp_dir().join(".env");
161        std::fs::write(&file, "BAR=BAZ\n\nexport QUUX = 'QUUUX'\n#IGNORE=ME").unwrap();
162
163        let mut base = HashMap::new();
164        base.insert("FOO".into(), "BAR".into());
165
166        let flags = vec!["FOO=QUX".to_string(), "BAZ=QUUX".to_string()];
167        let vars = lambda_environment(Some(&base), &Some(file), Some(&flags)).unwrap();
168
169        assert_eq!("QUX".to_string(), vars["FOO"]);
170        assert_eq!("QUUX".to_string(), vars["BAZ"]);
171        assert_eq!("BAZ".to_string(), vars["BAR"]);
172        assert_eq!("QUUUX".to_string(), vars["QUUX"]);
173        assert!(!vars.contains_key("IGNORE"));
174        assert!(!vars.contains_key(""));
175    }
176}