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, env::VarError, 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
110pub trait EnvVarExtractor {
111    fn var(&self, name: &str) -> Result<String, VarError>;
112}
113
114pub struct SystemEnvExtractor;
115
116impl EnvVarExtractor for SystemEnvExtractor {
117    fn var(&self, name: &str) -> Result<String, VarError> {
118        std::env::var(name)
119    }
120}
121
122pub struct HashMapEnvExtractor {
123    env: HashMap<String, String>,
124}
125
126impl From<Vec<(&str, &str)>> for HashMapEnvExtractor {
127    fn from(env: Vec<(&str, &str)>) -> Self {
128        Self {
129            env: env
130                .into_iter()
131                .map(|(k, v)| (k.to_string(), v.to_string()))
132                .collect(),
133        }
134    }
135}
136
137impl EnvVarExtractor for HashMapEnvExtractor {
138    fn var(&self, name: &str) -> Result<String, VarError> {
139        self.env.get(name).cloned().ok_or(VarError::NotPresent)
140    }
141}
142
143#[cfg(test)]
144mod test {
145    use std::env::temp_dir;
146
147    use super::*;
148
149    #[test]
150    fn test_extract_var() {
151        let (k, v) = extract_var("FOO=BAR").unwrap();
152        assert_eq!("FOO", k);
153        assert_eq!("BAR", v);
154
155        let (k, v) = extract_var(" FOO = BAR ").unwrap();
156        assert_eq!("FOO", k);
157        assert_eq!("BAR", v);
158
159        extract_var("=BAR").expect_err("missing key");
160        extract_var("FOO=").expect_err("missing value");
161        extract_var("  ").expect_err("missing variable");
162    }
163
164    #[test]
165    fn test_empty_environment() {
166        let env = lambda_environment(None, &None, None).unwrap();
167        assert!(env.is_empty());
168    }
169
170    #[test]
171    fn test_base_environment() {
172        let mut base = HashMap::new();
173        base.insert("FOO".into(), "BAR".into());
174        let env = lambda_environment(Some(&base), &None, None).unwrap();
175
176        assert_eq!("BAR".to_string(), env["FOO"]);
177    }
178
179    #[test]
180    fn test_environment_with_flags() {
181        let mut base = HashMap::new();
182        base.insert("FOO".into(), "BAR".into());
183
184        let flags = vec!["FOO=QUX".to_string(), "BAZ=QUUX".to_string()];
185        let env = lambda_environment(Some(&base), &None, Some(&flags)).unwrap();
186
187        assert_eq!("QUX".to_string(), env["FOO"]);
188        assert_eq!("QUUX".to_string(), env["BAZ"]);
189    }
190
191    #[test]
192    fn test_environment_with_file() {
193        let file = temp_dir().join(".env");
194        std::fs::write(&file, "BAR=BAZ\n\nexport QUUX = 'QUUUX'\n#IGNORE=ME").unwrap();
195
196        let mut base = HashMap::new();
197        base.insert("FOO".into(), "BAR".into());
198
199        let flags = vec!["FOO=QUX".to_string(), "BAZ=QUUX".to_string()];
200        let vars = lambda_environment(Some(&base), &Some(file), Some(&flags)).unwrap();
201
202        assert_eq!("QUX".to_string(), vars["FOO"]);
203        assert_eq!("QUUX".to_string(), vars["BAZ"]);
204        assert_eq!("BAZ".to_string(), vars["BAR"]);
205        assert_eq!("QUUUX".to_string(), vars["QUUX"]);
206        assert!(!vars.contains_key("IGNORE"));
207        assert!(!vars.contains_key(""));
208    }
209}
cargo_lambda_metadata/env.rs

cargo_lambda_metadata/
env.rs
