rtlambda/data/
env.rs

1// Copyright 2022 Guy Or and the "rtlambda" authors. All rights reserved.
2
3// `SPDX-License-Identifier: MIT OR Apache-2.0`
4
5/// An enum representing the `InitializationType` choices set as an env-var on the instance by AWS Lambda.
6/// See [Defined runtime environment variables](https://docs.aws.amazon.com/lambda/latest/dg/configuration-envvars.html#configuration-envvars-runtime).
7#[derive(Clone, Copy, Debug)]
8pub enum InitializationType {
9    OnDemand,
10    ProvisionedConcurrency,
11    Unknown,
12}
13
14impl InitializationType {
15    /// Returns the [`InitializationType`] value corresponding to the input string.
16    /// itype must be lowercase.
17    fn from_string(itype: &str) -> InitializationType {
18        match itype {
19            "on-demand" => Self::OnDemand,
20            "provisioned-concurrency" => Self::ProvisionedConcurrency,
21            // Shouldn't reach here but if for some reason AWS doesn't get it right...
22            _ => Self::Unknown,
23        }
24    }
25}
26
27/// An interface trait for reading the environment variables set by the AWS Lambda service.
28///
29/// Based on - [Defined runtime environment variables](https://docs.aws.amazon.com/lambda/latest/dg/configuration-envvars.html#configuration-envvars-runtime).
30pub trait RuntimeEnvVars: Default {
31    fn get_handler(&self) -> Option<&str>;
32    fn get_region(&self) -> Option<&str>;
33    fn get_trace_id(&self) -> Option<&str>;
34    fn get_execution_env(&self) -> Option<&str>;
35    fn get_function_name(&self) -> Option<&str>;
36    fn get_function_memory_size(&self) -> Option<usize>;
37    fn get_function_version(&self) -> Option<&str>;
38    fn get_initialization_type(&self) -> InitializationType;
39    fn get_log_group_name(&self) -> Option<&str>;
40    fn get_log_stream_name(&self) -> Option<&str>;
41    fn get_access_key(&self) -> Option<&str>;
42    fn get_access_key_id(&self) -> Option<&str>;
43    fn get_secret_access_key(&self) -> Option<&str>;
44    fn get_session_token(&self) -> Option<&str>;
45    fn get_runtime_api(&self) -> Option<&str>;
46    fn get_task_root(&self) -> Option<&str>;
47    fn get_runtime_dir(&self) -> Option<&str>;
48    fn get_tz(&self) -> Option<&str>;
49    /// Returns the string value of an env-var `var_name` wrapped in an [`Option`],
50    /// or `None` if the env-var is not set or the [`std::env::var`] function returns an error.
51    fn get_var(var_name: &str) -> Option<String> {
52        use std::env;
53        env::var(var_name).ok()
54    }
55    /// Signals that the previous tracing id has changed as a result of a new incoming event.
56    fn set_trace_id(&mut self, new_id: Option<&str>);
57}
58
59/// A struct implementing [`RuntimeEnvVars`] by caching the default runtime env-vars,
60/// and supports a default initialization using std::env::var calls.
61#[derive(Debug, Clone)]
62pub struct LambdaRuntimeEnv {
63    pub handler: Option<String>,
64    // This value should be set by the runtime after each next invocation request where a new id is given
65    pub trace_id: Option<String>,
66    pub region: Option<String>,
67    // Custom runtimes currently don't have this value set as per AWS docs
68    pub execution_env: Option<String>,
69    pub function_name: Option<String>,
70    pub function_memory_size: Option<usize>,
71    pub function_version: Option<String>,
72    pub initialization_type: InitializationType,
73    pub log_group_name: Option<String>,
74    pub log_stream_name: Option<String>,
75    pub access_key: Option<String>,
76    pub access_key_id: Option<String>,
77    pub secret_access_key: Option<String>,
78    pub session_token: Option<String>,
79    pub runtime_api: Option<String>,
80    pub task_root: Option<String>,
81    pub runtime_dir: Option<String>,
82    pub tz: Option<String>,
83}
84
85impl LambdaRuntimeEnv {
86    /// Constructs a new [`LambdaRuntimeEnv`] by reading the process' environment variables,
87    /// and caching the [default env-vars](https://docs.aws.amazon.com/lambda/latest/dg/configuration-envvars.html#configuration-envvars-runtime).
88    pub fn from_env() -> LambdaRuntimeEnv {
89        use std::env;
90        LambdaRuntimeEnv {
91            handler: env::var("_HANDLER").ok(),
92            region: env::var("AWS_REGION").ok(),
93            trace_id: None,
94            execution_env: env::var("AWS_EXECUTION_ENV").ok(),
95            function_name: env::var("AWS_LAMBDA_FUNCTION_NAME").ok(),
96            function_memory_size: match env::var("AWS_LAMBDA_FUNCTION_MEMORY_SIZE").ok() {
97                Some(v) => v.parse::<usize>().ok(),
98                None => None,
99            },
100            function_version: env::var("AWS_LAMBDA_FUNCTION_VERSION").ok(),
101            initialization_type: match env::var("AWS_LAMBDA_INITIALIZATION_TYPE").ok() {
102                Some(v) => InitializationType::from_string(&v),
103                None => InitializationType::Unknown,
104            },
105            log_group_name: env::var("AWS_LAMBDA_LOG_GROUP_NAME").ok(),
106            log_stream_name: env::var("AWS_LAMBDA_LOG_STREAM_NAME").ok(),
107            access_key: env::var("AWS_ACCESS_KEY").ok(),
108            access_key_id: env::var("AWS_ACCESS_KEY_ID").ok(),
109            secret_access_key: env::var("AWS_SECRET_ACCESS_KEY").ok(),
110            session_token: env::var("AWS_SESSION_TOKEN").ok(),
111            runtime_api: env::var("AWS_LAMBDA_RUNTIME_API").ok(),
112            task_root: env::var("LAMBDA_TASK_ROOT").ok(),
113            runtime_dir: env::var("LAMBDA_RUNTIME_DIR").ok(),
114            tz: env::var("TZ").ok(),
115        }
116    }
117}
118
119impl Default for LambdaRuntimeEnv {
120    fn default() -> Self {
121        Self::from_env()
122    }
123}
124
125impl RuntimeEnvVars for LambdaRuntimeEnv {
126    #[inline(always)]
127    fn get_handler(&self) -> Option<&str> {
128        self.handler.as_deref()
129    }
130
131    #[inline(always)]
132    fn get_region(&self) -> Option<&str> {
133        self.region.as_deref()
134    }
135
136    #[inline(always)]
137    fn get_trace_id(&self) -> Option<&str> {
138        self.trace_id.as_deref()
139    }
140
141    #[inline(always)]
142    fn get_execution_env(&self) -> Option<&str> {
143        self.execution_env.as_deref()
144    }
145
146    #[inline(always)]
147    fn get_function_name(&self) -> Option<&str> {
148        self.function_name.as_deref()
149    }
150
151    #[inline(always)]
152    fn get_function_memory_size(&self) -> Option<usize> {
153        self.function_memory_size
154    }
155
156    #[inline(always)]
157    fn get_function_version(&self) -> Option<&str> {
158        self.function_version.as_deref()
159    }
160
161    #[inline(always)]
162    fn get_initialization_type(&self) -> InitializationType {
163        self.initialization_type
164    }
165    #[inline(always)]
166    fn get_log_group_name(&self) -> Option<&str> {
167        self.log_group_name.as_deref()
168    }
169
170    #[inline(always)]
171    fn get_log_stream_name(&self) -> Option<&str> {
172        self.log_stream_name.as_deref()
173    }
174
175    #[inline(always)]
176    fn get_access_key(&self) -> Option<&str> {
177        self.access_key.as_deref()
178    }
179
180    #[inline(always)]
181    fn get_access_key_id(&self) -> Option<&str> {
182        self.access_key_id.as_deref()
183    }
184
185    #[inline(always)]
186    fn get_secret_access_key(&self) -> Option<&str> {
187        self.secret_access_key.as_deref()
188    }
189
190    #[inline(always)]
191    fn get_session_token(&self) -> Option<&str> {
192        self.session_token.as_deref()
193    }
194
195    #[inline(always)]
196    fn get_runtime_api(&self) -> Option<&str> {
197        self.runtime_api.as_deref()
198    }
199
200    #[inline(always)]
201    fn get_task_root(&self) -> Option<&str> {
202        self.task_root.as_deref()
203    }
204
205    #[inline(always)]
206    fn get_runtime_dir(&self) -> Option<&str> {
207        self.runtime_dir.as_deref()
208    }
209
210    #[inline(always)]
211    fn get_tz(&self) -> Option<&str> {
212        self.tz.as_deref()
213    }
214
215    #[inline]
216    fn set_trace_id(&mut self, new_id: Option<&str>) {
217        self.trace_id = new_id.map(|v| v.to_string());
218    }
219}