Skip to main content

lambda_forge/
endpoint.rs

1use crate::Error;
2
3use lambda_http::Context;
4use or_panic::prelude::*;
5
6#[derive(Debug, Clone)]
7/// Metadata for a specific endpoint.
8pub struct EndpointMetadata {
9    /// Product name (Application name).
10    pub product: String,
11
12    /// Component or section name.
13    pub component: String,
14
15    /// The AWS Lambda function execution context.
16    ///
17    /// NOTE: The values in this struct are populated using
18    /// the Lambda environment variables and the headers
19    /// returned by the poll request to the Runtime APIs.
20    pub lambda_context: Context,
21
22    /// Current stage it's running within.
23    pub stage: Option<String>,
24
25    /// Response content type for the endpoint.
26    pub response_content_type: String,
27}
28impl EndpointMetadata {
29    pub fn builder() -> EndpointMetadataBuilder {
30        EndpointMetadataBuilder::default()
31    }
32
33    /// Get the URL to the lambda function currently running.
34    pub fn lambda_url(&self) -> Result<String, Error> {
35        let fn_name = &self.lambda_context.env_config.function_name;
36        let fn_region = std::env::var("AWS_REGION")
37            .or_panic("AWS Lambda should always have the AWS_REGION environment variable set.");
38
39        let mut url = format!(
40            "https://{}.console.aws.amazon.com/lambda/home?region={}#/functions/{}?tab=monitoring",
41            fn_region, fn_region, fn_name
42        );
43
44        if let Some(stage) = &self.stage {
45            url.push_str(&format!("&qualifier={stage}"));
46        }
47
48        Ok(url)
49    }
50
51    /// Get the URL to the cloudwatch log of the invocation currently running.
52    pub fn cloudwatch_log(&self, request_id: &str) -> Result<String, Error> {
53        let fn_name = &self.lambda_context.env_config.function_name;
54        let fn_region = std::env::var("AWS_REGION")
55            .or_panic("AWS Lambda should always have the AWS_REGION environment variable set.");
56
57        let url = format!(
58            concat!(
59                "https://{fn_region}.console.aws.amazon.com/",
60                "cloudwatch/home?region={fn_region}#logsV2:log-groups/log-group/",
61                "$252Faws$252Flambda$252F{fn_name}/log-events",
62                "$3FfilterPattern$3D$2522{request_id}$2522"
63            ),
64            fn_region = fn_region,
65            fn_name = fn_name,
66            request_id = request_id,
67        );
68
69        Ok(url)
70    }
71}
72
73#[derive(Default)]
74pub struct EndpointMetadataBuilder {
75    pub product: Option<String>,
76    pub component: Option<String>,
77    pub lambda_context: Option<Context>,
78    pub stage: Option<String>,
79    pub response_content_type: Option<String>,
80}
81impl EndpointMetadataBuilder {
82    pub fn new() -> Self {
83        Self::default()
84    }
85
86    pub fn context(mut self, ctx: Context) -> Self {
87        let parts: Vec<&str> = ctx.invoked_function_arn.split(':').collect();
88        let alias = parts.get(7).map(|s| (*s).to_owned());
89
90        self.stage = alias;
91        self.lambda_context = Some(ctx);
92        self
93    }
94
95    pub fn product(mut self, v: impl Into<String>) -> Self {
96        self.product = Some(v.into());
97        self
98    }
99
100    pub fn component(mut self, v: impl Into<String>) -> Self {
101        self.component = Some(v.into());
102        self
103    }
104
105    pub fn stage(mut self, v: impl Into<String>) -> Self {
106        self.stage = Some(v.into());
107        self
108    }
109
110    pub fn response_content_type(mut self, v: impl Into<String>) -> Self {
111        self.response_content_type = Some(v.into());
112        self
113    }
114
115    pub fn build(self) -> Result<EndpointMetadata, Error> {
116        let aws_lambda_context = self.lambda_context.ok_or(Error::MissingLambdaContext)?;
117
118        Ok(EndpointMetadata {
119            product: self.product.unwrap_or("N/A".to_owned()),
120            component: self.component.unwrap_or("N/A".to_owned()),
121            lambda_context: aws_lambda_context,
122            stage: self.stage,
123            response_content_type: self
124                .response_content_type
125                .unwrap_or("application/json".to_owned()),
126        })
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use super::*;
133    use lambda_http::lambda_runtime::Config;
134    use std::sync::{Arc, Mutex, OnceLock};
135
136    fn env_lock() -> &'static Mutex<()> {
137        static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
138        LOCK.get_or_init(|| Mutex::new(()))
139    }
140
141    fn with_aws_region<T>(region: &str, f: impl FnOnce() -> T) -> T {
142        let _g = env_lock().lock().unwrap();
143
144        let prev = std::env::var("AWS_REGION").ok();
145        unsafe { std::env::set_var("AWS_REGION", region) };
146
147        let out = f();
148
149        match prev {
150            Some(v) => unsafe { std::env::set_var("AWS_REGION", v) },
151            None => unsafe { std::env::remove_var("AWS_REGION") },
152        }
153        out
154    }
155
156    fn mk_ctx(function_name: &str, invoked_arn: &str) -> Context {
157        let mut ctx = Context::default();
158        ctx.invoked_function_arn = invoked_arn.to_string();
159
160        let cfg = Config {
161            function_name: function_name.to_string(),
162            ..Default::default()
163        };
164        ctx.env_config = Arc::new(cfg);
165
166        ctx
167    }
168
169    #[test]
170    fn builder_build_requires_context() {
171        let err = EndpointMetadata::builder().build().unwrap_err();
172        match err {
173            Error::MissingLambdaContext => {}
174            other => panic!("expected MissingLambdaContext, got: {other:?}"),
175        }
176    }
177
178    #[test]
179    fn builder_defaults_when_fields_missing() {
180        let ctx = mk_ctx(
181            "my-fn",
182            "arn:aws:lambda:us-east-1:123456789012:function:my-fn",
183        );
184        let md = EndpointMetadata::builder().context(ctx).build().unwrap();
185
186        assert_eq!(md.product, "N/A");
187        assert_eq!(md.component, "N/A");
188        assert_eq!(md.response_content_type, "application/json");
189        assert_eq!(md.stage, None);
190    }
191
192    #[test]
193    fn builder_context_extracts_stage_from_invoked_arn_alias() {
194        let arn = "arn:aws:lambda:us-east-1:123456789012:function:my-fn:dev";
195        let ctx = mk_ctx("my-fn", arn);
196
197        let md = EndpointMetadata::builder().context(ctx).build().unwrap();
198        assert_eq!(md.stage.as_deref(), Some("dev"));
199    }
200
201    #[test]
202    fn builder_context_no_alias_keeps_stage_none() {
203        let arn = "arn:aws:lambda:us-east-1:123456789012:function:my-fn";
204        let ctx = mk_ctx("my-fn", arn);
205
206        let md = EndpointMetadata::builder().context(ctx).build().unwrap();
207        assert_eq!(md.stage, None);
208    }
209
210    #[test]
211    fn lambda_url_without_stage_has_no_qualifier() {
212        let ctx = mk_ctx("my-fn", "arn:aws:lambda:us-east-1:123:function:my-fn");
213        let md = EndpointMetadata::builder().context(ctx).build().unwrap();
214
215        let url = with_aws_region("us-east-1", || md.lambda_url().unwrap());
216
217        assert_eq!(
218            url,
219            "https://us-east-1.console.aws.amazon.com/lambda/home?region=us-east-1#/functions/my-fn?tab=monitoring"
220        );
221        assert!(!url.contains("qualifier="));
222    }
223
224    #[test]
225    fn lambda_url_with_stage_appends_qualifier() {
226        let ctx = mk_ctx("my-fn", "arn:aws:lambda:us-east-1:123:function:my-fn");
227        let md = EndpointMetadata::builder()
228            .context(ctx)
229            .stage("prod")
230            .build()
231            .unwrap();
232
233        let url = with_aws_region("us-east-1", || md.lambda_url().unwrap());
234
235        assert_eq!(
236            url,
237            "https://us-east-1.console.aws.amazon.com/lambda/home?region=us-east-1#/functions/my-fn?tab=monitoring&qualifier=prod"
238        );
239    }
240
241    #[test]
242    fn cloudwatch_log_formats_expected_url() {
243        let ctx = mk_ctx("my-fn", "arn:aws:lambda:us-east-1:123:function:my-fn");
244        let md = EndpointMetadata::builder().context(ctx).build().unwrap();
245
246        let url = with_aws_region("us-west-2", || md.cloudwatch_log("req-123").unwrap());
247
248        assert_eq!(
249            url,
250            "https://us-west-2.console.aws.amazon.com/cloudwatch/home?region=us-west-2#logsV2:log-groups/log-group/$252Faws$252Flambda$252Fmy-fn/log-events$3FfilterPattern$3D$2522req-123$2522"
251        );
252    }
253
254    #[test]
255    fn builder_setters_override_defaults() {
256        let ctx = mk_ctx("my-fn", "arn:aws:lambda:us-east-1:123:function:my-fn");
257
258        let md = EndpointMetadata::builder()
259            .context(ctx)
260            .product("Couplet")
261            .component("Auth")
262            .response_content_type("text/plain")
263            .build()
264            .unwrap();
265
266        assert_eq!(md.product, "Couplet");
267        assert_eq!(md.component, "Auth");
268        assert_eq!(md.response_content_type, "text/plain");
269    }
270}