lambda-forge 0.1.7

use crate::Error;

use lambda_http::Context;
use or_panic::prelude::*;

#[derive(Debug, Clone)]
/// Metadata for a specific endpoint.
pub struct EndpointMetadata {
    /// Product name (Application name).
    pub product: String,

    /// Component or section name.
    pub component: String,

    /// The AWS Lambda function execution context.
    ///
    /// NOTE: The values in this struct are populated using
    /// the Lambda environment variables and the headers
    /// returned by the poll request to the Runtime APIs.
    pub lambda_context: Context,

    /// Current stage it's running within.
    pub stage: Option<String>,

    /// Response content type for the endpoint.
    pub response_content_type: String,
}
impl EndpointMetadata {
    pub fn builder() -> EndpointMetadataBuilder {
        EndpointMetadataBuilder::default()
    }

    /// Get the URL to the lambda function currently running.
    pub fn lambda_url(&self) -> Result<String, Error> {
        let fn_name = &self.lambda_context.env_config.function_name;
        let fn_region = std::env::var("AWS_REGION")
            .or_panic("AWS Lambda should always have the AWS_REGION environment variable set.");

        let mut url = format!(
            "https://{}.console.aws.amazon.com/lambda/home?region={}#/functions/{}?tab=monitoring",
            fn_region, fn_region, fn_name
        );

        if let Some(stage) = &self.stage {
            url.push_str(&format!("&qualifier={stage}"));
        }

        Ok(url)
    }

    /// Get the URL to the cloudwatch log of the invocation currently running.
    pub fn cloudwatch_log(&self, request_id: &str) -> Result<String, Error> {
        let fn_name = &self.lambda_context.env_config.function_name;
        let fn_region = std::env::var("AWS_REGION")
            .or_panic("AWS Lambda should always have the AWS_REGION environment variable set.");

        let url = format!(
            concat!(
                "https://{fn_region}.console.aws.amazon.com/",
                "cloudwatch/home?region={fn_region}#logsV2:log-groups/log-group/",
                "$252Faws$252Flambda$252F{fn_name}/log-events",
                "$3FfilterPattern$3D$2522{request_id}$2522"
            ),
            fn_region = fn_region,
            fn_name = fn_name,
            request_id = request_id,
        );

        Ok(url)
    }
}

#[derive(Default)]
pub struct EndpointMetadataBuilder {
    pub product: Option<String>,
    pub component: Option<String>,
    pub lambda_context: Option<Context>,
    pub stage: Option<String>,
    pub response_content_type: Option<String>,
}
impl EndpointMetadataBuilder {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn context(mut self, ctx: Context) -> Self {
        let parts: Vec<&str> = ctx.invoked_function_arn.split(':').collect();
        let alias = parts.get(7).map(|s| (*s).to_owned());

        self.stage = alias;
        self.lambda_context = Some(ctx);
        self
    }

    pub fn product(mut self, v: impl Into<String>) -> Self {
        self.product = Some(v.into());
        self
    }

    pub fn component(mut self, v: impl Into<String>) -> Self {
        self.component = Some(v.into());
        self
    }

    pub fn stage(mut self, v: impl Into<String>) -> Self {
        self.stage = Some(v.into());
        self
    }

    pub fn response_content_type(mut self, v: impl Into<String>) -> Self {
        self.response_content_type = Some(v.into());
        self
    }

    pub fn build(self) -> Result<EndpointMetadata, Error> {
        let aws_lambda_context = self.lambda_context.ok_or(Error::MissingLambdaContext)?;

        Ok(EndpointMetadata {
            product: self.product.unwrap_or("N/A".to_owned()),
            component: self.component.unwrap_or("N/A".to_owned()),
            lambda_context: aws_lambda_context,
            stage: self.stage,
            response_content_type: self
                .response_content_type
                .unwrap_or("application/json".to_owned()),
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use lambda_http::lambda_runtime::Config;
    use std::sync::{Arc, Mutex, OnceLock};

    fn env_lock() -> &'static Mutex<()> {
        static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
        LOCK.get_or_init(|| Mutex::new(()))
    }

    fn with_aws_region<T>(region: &str, f: impl FnOnce() -> T) -> T {
        let _g = env_lock().lock().unwrap();

        let prev = std::env::var("AWS_REGION").ok();
        unsafe { std::env::set_var("AWS_REGION", region) };

        let out = f();

        match prev {
            Some(v) => unsafe { std::env::set_var("AWS_REGION", v) },
            None => unsafe { std::env::remove_var("AWS_REGION") },
        }
        out
    }

    fn mk_ctx(function_name: &str, invoked_arn: &str) -> Context {
        let mut ctx = Context::default();
        ctx.invoked_function_arn = invoked_arn.to_string();

        let cfg = Config {
            function_name: function_name.to_string(),
            ..Default::default()
        };
        ctx.env_config = Arc::new(cfg);

        ctx
    }

    #[test]
    fn builder_build_requires_context() {
        let err = EndpointMetadata::builder().build().unwrap_err();
        match err {
            Error::MissingLambdaContext => {}
            other => panic!("expected MissingLambdaContext, got: {other:?}"),
        }
    }

    #[test]
    fn builder_defaults_when_fields_missing() {
        let ctx = mk_ctx(
            "my-fn",
            "arn:aws:lambda:us-east-1:123456789012:function:my-fn",
        );
        let md = EndpointMetadata::builder().context(ctx).build().unwrap();

        assert_eq!(md.product, "N/A");
        assert_eq!(md.component, "N/A");
        assert_eq!(md.response_content_type, "application/json");
        assert_eq!(md.stage, None);
    }

    #[test]
    fn builder_context_extracts_stage_from_invoked_arn_alias() {
        let arn = "arn:aws:lambda:us-east-1:123456789012:function:my-fn:dev";
        let ctx = mk_ctx("my-fn", arn);

        let md = EndpointMetadata::builder().context(ctx).build().unwrap();
        assert_eq!(md.stage.as_deref(), Some("dev"));
    }

    #[test]
    fn builder_context_no_alias_keeps_stage_none() {
        let arn = "arn:aws:lambda:us-east-1:123456789012:function:my-fn";
        let ctx = mk_ctx("my-fn", arn);

        let md = EndpointMetadata::builder().context(ctx).build().unwrap();
        assert_eq!(md.stage, None);
    }

    #[test]
    fn lambda_url_without_stage_has_no_qualifier() {
        let ctx = mk_ctx("my-fn", "arn:aws:lambda:us-east-1:123:function:my-fn");
        let md = EndpointMetadata::builder().context(ctx).build().unwrap();

        let url = with_aws_region("us-east-1", || md.lambda_url().unwrap());

        assert_eq!(
            url,
            "https://us-east-1.console.aws.amazon.com/lambda/home?region=us-east-1#/functions/my-fn?tab=monitoring"
        );
        assert!(!url.contains("qualifier="));
    }

    #[test]
    fn lambda_url_with_stage_appends_qualifier() {
        let ctx = mk_ctx("my-fn", "arn:aws:lambda:us-east-1:123:function:my-fn");
        let md = EndpointMetadata::builder()
            .context(ctx)
            .stage("prod")
            .build()
            .unwrap();

        let url = with_aws_region("us-east-1", || md.lambda_url().unwrap());

        assert_eq!(
            url,
            "https://us-east-1.console.aws.amazon.com/lambda/home?region=us-east-1#/functions/my-fn?tab=monitoring&qualifier=prod"
        );
    }

    #[test]
    fn cloudwatch_log_formats_expected_url() {
        let ctx = mk_ctx("my-fn", "arn:aws:lambda:us-east-1:123:function:my-fn");
        let md = EndpointMetadata::builder().context(ctx).build().unwrap();

        let url = with_aws_region("us-west-2", || md.cloudwatch_log("req-123").unwrap());

        assert_eq!(
            url,
            "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"
        );
    }

    #[test]
    fn builder_setters_override_defaults() {
        let ctx = mk_ctx("my-fn", "arn:aws:lambda:us-east-1:123:function:my-fn");

        let md = EndpointMetadata::builder()
            .context(ctx)
            .product("Couplet")
            .component("Auth")
            .response_content_type("text/plain")
            .build()
            .unwrap();

        assert_eq!(md.product, "Couplet");
        assert_eq!(md.component, "Auth");
        assert_eq!(md.response_content_type, "text/plain");
    }
}