diff --git a/.github/workflows/run-integration-test.yml b/.github/workflows/run-integration-test.yml index a4fd604b..523aadfd 100644 --- a/.github/workflows/run-integration-test.yml +++ b/.github/workflows/run-integration-test.yml @@ -47,12 +47,15 @@ jobs: echo "STACK_NAME=$stackName" >> "$GITHUB_OUTPUT" echo "Stack name = $stackName" sam deploy --stack-name "${stackName}" --parameter-overrides "ParameterKey=SecretToken,ParameterValue=${{ secrets.SECRET_TOKEN }}" "ParameterKey=LambdaRole,ParameterValue=${{ secrets.AWS_LAMBDA_ROLE }}" --no-confirm-changeset --no-progressbar > disable_output - TEST_ENDPOINT=$(sam list stack-outputs --stack-name "${stackName}" --output json | jq -r '.[] | .OutputValue') + TEST_ENDPOINT=$(sam list stack-outputs --stack-name "${stackName}" --output json | jq -r '.[] | select(.OutputKey=="HelloApiEndpoint") | .OutputValue') + TENANT_ID_TEST_FUNCTION=$(sam list stack-outputs --stack-name "${stackName}" --output json | jq -r '.[] | select(.OutputKey=="TenantIdTestFunction") | .OutputValue') echo "TEST_ENDPOINT=$TEST_ENDPOINT" >> "$GITHUB_OUTPUT" + echo "TENANT_ID_TEST_FUNCTION=$TENANT_ID_TEST_FUNCTION" >> "$GITHUB_OUTPUT" - name: run test env: SECRET_TOKEN: ${{ secrets.SECRET_TOKEN }} TEST_ENDPOINT: ${{ steps.deploy_stack.outputs.TEST_ENDPOINT }} + TENANT_ID_TEST_FUNCTION: ${{ steps.deploy_stack.outputs.TENANT_ID_TEST_FUNCTION }} run: cd lambda-integration-tests && cargo test - name: cleanup if: always() diff --git a/examples/basic-tenant-id/Cargo.toml b/examples/basic-tenant-id/Cargo.toml new file mode 100644 index 00000000..2b1144cb --- /dev/null +++ b/examples/basic-tenant-id/Cargo.toml @@ -0,0 +1,11 @@ +[package] +name = "basic-tenant-id" +version = "0.1.0" +edition = "2021" + +[dependencies] +lambda_runtime = { path = "../../lambda-runtime" } +serde_json = "1.0" +tokio = { version = "1", features = ["macros"] } +tracing = { version = "0.1", features = ["log"] } +tracing-subscriber = { version = "0.3", default-features = false, features = ["fmt"] } diff --git a/examples/basic-tenant-id/README.md b/examples/basic-tenant-id/README.md new file mode 100644 index 00000000..eb233e75 --- /dev/null +++ b/examples/basic-tenant-id/README.md @@ -0,0 +1,38 @@ +# Basic Tenant ID Example + +This example demonstrates how to access and use tenant ID information in a Lambda function. + +## Key Features + +- Extracts tenant ID from Lambda runtime headers +- Includes tenant ID in tracing logs +- Returns tenant ID in the response +- Handles cases where tenant ID is not provided + +## Usage + +The tenant ID is automatically extracted from the `lambda-runtime-aws-tenant-id` header and made available in the Lambda context. + +```rust +async fn function_handler(event: LambdaEvent) -> Result { + let (event, context) = event.into_parts(); + + // Access tenant ID from context + match &context.tenant_id { + Some(tenant_id) => println!("Processing for tenant: {}", tenant_id), + None => println!("No tenant ID provided"), + } + + // ... rest of function logic +} +``` + +## Testing + +You can test this function locally using cargo lambda: + +```bash +cargo lambda invoke --data-ascii '{"test": "data"}' +``` + +The tenant ID will be None when testing locally unless you set up a mock runtime environment with the appropriate headers. diff --git a/examples/basic-tenant-id/src/main.rs b/examples/basic-tenant-id/src/main.rs new file mode 100644 index 00000000..1789bcba --- /dev/null +++ b/examples/basic-tenant-id/src/main.rs @@ -0,0 +1,36 @@ +use lambda_runtime::{service_fn, Error, LambdaEvent}; +use serde_json::{json, Value}; + +async fn function_handler(event: LambdaEvent) -> Result { + let (event, context) = event.into_parts(); + + // Access tenant ID from context + let tenant_info = match &context.tenant_id { + Some(tenant_id) => format!("Processing request for tenant: {}", tenant_id), + None => "No tenant ID provided".to_string(), + }; + + tracing::info!("Request ID: {}", context.request_id); + tracing::info!("Tenant info: {}", tenant_info); + + // Include tenant ID in response + let response = json!({ + "message": "Hello from Lambda!", + "request_id": context.request_id, + "tenant_id": context.tenant_id, + "input": event + }); + + Ok(response) +} + +#[tokio::main] +async fn main() -> Result<(), Error> { + tracing_subscriber::fmt() + .with_max_level(tracing::Level::INFO) + .with_target(false) + .without_time() + .init(); + + lambda_runtime::run(service_fn(function_handler)).await +} diff --git a/lambda-integration-tests/Cargo.toml b/lambda-integration-tests/Cargo.toml index 9fc7aff5..b7b505d9 100644 --- a/lambda-integration-tests/Cargo.toml +++ b/lambda-integration-tests/Cargo.toml @@ -17,6 +17,7 @@ aws_lambda_events = { path = "../lambda-events" } serde_json = "1.0.121" tokio = { version = "1", features = ["full"] } serde = { version = "1.0.204", features = ["derive"] } +tracing = "0.1" [dev-dependencies] reqwest = { version = "0.12.5", features = ["blocking"] } @@ -31,3 +32,7 @@ path = "src/helloworld.rs" [[bin]] name = "authorizer" path = "src/authorizer.rs" + +[[bin]] +name = "tenant-id-test" +path = "src/tenant_id_test.rs" diff --git a/lambda-integration-tests/src/tenant_id_test.rs b/lambda-integration-tests/src/tenant_id_test.rs new file mode 100644 index 00000000..aeebc182 --- /dev/null +++ b/lambda-integration-tests/src/tenant_id_test.rs @@ -0,0 +1,26 @@ +use lambda_runtime::{service_fn, Error, LambdaEvent}; +use serde_json::{json, Value}; + +async fn function_handler(event: LambdaEvent) -> Result { + let (event, context) = event.into_parts(); + + tracing::info!("Processing request with tenant ID: {:?}", context.tenant_id); + + let response = json!({ + "statusCode": 200, + "body": json!({ + "message": "Tenant ID test successful", + "request_id": context.request_id, + "tenant_id": context.tenant_id, + "input": event + }).to_string() + }); + + Ok(response) +} + +#[tokio::main] +async fn main() -> Result<(), Error> { + lambda_runtime::tracing::init_default_subscriber(); + lambda_runtime::run(service_fn(function_handler)).await +} diff --git a/lambda-integration-tests/template.yaml b/lambda-integration-tests/template.yaml index 1aa69fc8..9ba694e3 100644 --- a/lambda-integration-tests/template.yaml +++ b/lambda-integration-tests/template.yaml @@ -41,6 +41,18 @@ Resources: Path: /hello Method: get + TenantIdTestFunction: + Type: AWS::Serverless::Function + Metadata: + BuildMethod: rust-cargolambda + BuildProperties: + Binary: tenant-id-test + Properties: + CodeUri: ./ + Handler: bootstrap + Runtime: provided.al2023 + Role: !Ref LambdaRole + AuthorizerFunction: Type: AWS::Serverless::Function Metadata: @@ -59,4 +71,7 @@ Resources: Outputs: HelloApiEndpoint: Description: "API Gateway endpoint URL for HelloWorld" - Value: !Sub "https://${API}.execute-api.${AWS::Region}.amazonaws.com/integ-test/hello/" \ No newline at end of file + Value: !Sub "https://${API}.execute-api.${AWS::Region}.amazonaws.com/integ-test/hello/" + TenantIdTestFunction: + Description: "Tenant ID test function name" + Value: !Ref TenantIdTestFunction \ No newline at end of file diff --git a/lambda-runtime/src/layers/mod.rs b/lambda-runtime/src/layers/mod.rs index a05b6c67..22db335a 100644 --- a/lambda-runtime/src/layers/mod.rs +++ b/lambda-runtime/src/layers/mod.rs @@ -4,7 +4,7 @@ mod api_response; mod panic; // Publicly available services. -mod trace; +pub mod trace; pub(crate) use api_client::RuntimeApiClientService; pub(crate) use api_response::RuntimeApiResponseService; diff --git a/lambda-runtime/src/layers/otel.rs b/lambda-runtime/src/layers/otel.rs index 5e96dfed..3704b221 100644 --- a/lambda-runtime/src/layers/otel.rs +++ b/lambda-runtime/src/layers/otel.rs @@ -72,14 +72,26 @@ where } fn call(&mut self, req: LambdaInvocation) -> Self::Future { - let span = tracing::info_span!( - "Lambda function invocation", - "otel.name" = req.context.env_config.function_name, - "otel.kind" = field::Empty, - { attribute::FAAS_TRIGGER } = &self.otel_attribute_trigger, - { attribute::FAAS_INVOCATION_ID } = req.context.request_id, - { attribute::FAAS_COLDSTART } = self.coldstart - ); + let span = if let Some(tenant_id) = &req.context.tenant_id { + tracing::info_span!( + "Lambda function invocation", + "otel.name" = req.context.env_config.function_name, + "otel.kind" = field::Empty, + { attribute::FAAS_TRIGGER } = &self.otel_attribute_trigger, + { attribute::FAAS_INVOCATION_ID } = req.context.request_id, + { attribute::FAAS_COLDSTART } = self.coldstart, + "tenant_id" = tenant_id + ) + } else { + tracing::info_span!( + "Lambda function invocation", + "otel.name" = req.context.env_config.function_name, + "otel.kind" = field::Empty, + { attribute::FAAS_TRIGGER } = &self.otel_attribute_trigger, + { attribute::FAAS_INVOCATION_ID } = req.context.request_id, + { attribute::FAAS_COLDSTART } = self.coldstart + ) + }; // After the first execution, we can set 'coldstart' to false self.coldstart = false; diff --git a/lambda-runtime/src/layers/trace.rs b/lambda-runtime/src/layers/trace.rs index e93927b1..00de7781 100644 --- a/lambda-runtime/src/layers/trace.rs +++ b/lambda-runtime/src/layers/trace.rs @@ -55,16 +55,31 @@ where /* ------------------------------------------- UTILS ------------------------------------------- */ -fn request_span(ctx: &Context) -> tracing::Span { - match &ctx.xray_trace_id { - Some(trace_id) => { +pub fn request_span(ctx: &Context) -> tracing::Span { + match (&ctx.xray_trace_id, &ctx.tenant_id) { + (Some(trace_id), Some(tenant_id)) => { + tracing::info_span!( + "Lambda runtime invoke", + requestId = &ctx.request_id, + xrayTraceId = trace_id, + tenantId = tenant_id + ) + } + (Some(trace_id), None) => { tracing::info_span!( "Lambda runtime invoke", requestId = &ctx.request_id, xrayTraceId = trace_id ) } - None => { + (None, Some(tenant_id)) => { + tracing::info_span!( + "Lambda runtime invoke", + requestId = &ctx.request_id, + tenantId = tenant_id + ) + } + (None, None) => { tracing::info_span!("Lambda runtime invoke", requestId = &ctx.request_id) } } diff --git a/lambda-runtime/src/types.rs b/lambda-runtime/src/types.rs index 5e5f487a..8be45b59 100644 --- a/lambda-runtime/src/types.rs +++ b/lambda-runtime/src/types.rs @@ -79,6 +79,8 @@ pub struct Context { /// unless the invocation request to the Lambda APIs was made using AWS /// credentials issues by Amazon Cognito Identity Pools. pub identity: Option, + /// The tenant ID for the current invocation. + pub tenant_id: Option, /// Lambda function configuration from the local environment variables. /// Includes information such as the function name, memory allocation, /// version, and log streams. @@ -94,6 +96,7 @@ impl Default for Context { xray_trace_id: None, client_context: None, identity: None, + tenant_id: None, env_config: std::sync::Arc::new(crate::Config::default()), } } @@ -134,6 +137,9 @@ impl Context { .map(|v| String::from_utf8_lossy(v.as_bytes()).to_string()), client_context, identity, + tenant_id: headers + .get("lambda-runtime-aws-tenant-id") + .map(|v| String::from_utf8_lossy(v.as_bytes()).to_string()), env_config, }; @@ -496,4 +502,27 @@ mod test { assert_eq!(metadata_prelude, deserialized); } + + #[test] + fn context_with_tenant_id_resolves() { + let config = Arc::new(Config::default()); + let mut headers = HeaderMap::new(); + headers.insert("lambda-runtime-aws-request-id", HeaderValue::from_static("my-id")); + headers.insert("lambda-runtime-deadline-ms", HeaderValue::from_static("123")); + headers.insert("lambda-runtime-aws-tenant-id", HeaderValue::from_static("tenant-123")); + + let context = Context::new("id", config, &headers).unwrap(); + assert_eq!(context.tenant_id, Some("tenant-123".to_string())); + } + + #[test] + fn context_without_tenant_id_resolves() { + let config = Arc::new(Config::default()); + let mut headers = HeaderMap::new(); + headers.insert("lambda-runtime-aws-request-id", HeaderValue::from_static("my-id")); + headers.insert("lambda-runtime-deadline-ms", HeaderValue::from_static("123")); + + let context = Context::new("id", config, &headers).unwrap(); + assert_eq!(context.tenant_id, None); + } } diff --git a/lambda-runtime/tests/tenant_id_integration.rs b/lambda-runtime/tests/tenant_id_integration.rs new file mode 100644 index 00000000..39d65ccf --- /dev/null +++ b/lambda-runtime/tests/tenant_id_integration.rs @@ -0,0 +1,93 @@ +use http::HeaderMap; +use lambda_runtime::Context; +use std::sync::Arc; + +#[test] +fn test_context_tenant_id_extraction() { + let config = Arc::new(lambda_runtime::Config::default()); + + // Test with tenant ID + let mut headers = HeaderMap::new(); + headers.insert("lambda-runtime-aws-request-id", "test-id".parse().unwrap()); + headers.insert("lambda-runtime-deadline-ms", "123456789".parse().unwrap()); + headers.insert("lambda-runtime-aws-tenant-id", "my-tenant-123".parse().unwrap()); + + let context = Context::new("test-id", config.clone(), &headers).unwrap(); + assert_eq!(context.tenant_id, Some("my-tenant-123".to_string())); + + // Test without tenant ID + let mut headers = HeaderMap::new(); + headers.insert("lambda-runtime-aws-request-id", "test-id".parse().unwrap()); + headers.insert("lambda-runtime-deadline-ms", "123456789".parse().unwrap()); + + let context = Context::new("test-id", config, &headers).unwrap(); + assert_eq!(context.tenant_id, None); +} + +#[test] +fn test_context_tenant_id_with_special_characters() { + let config = Arc::new(lambda_runtime::Config::default()); + + // Test with tenant ID containing special characters + let mut headers = HeaderMap::new(); + headers.insert("lambda-runtime-aws-request-id", "test-id".parse().unwrap()); + headers.insert("lambda-runtime-deadline-ms", "123456789".parse().unwrap()); + headers.insert( + "lambda-runtime-aws-tenant-id", + "tenant-with-dashes_and_underscores.123".parse().unwrap(), + ); + + let context = Context::new("test-id", config, &headers).unwrap(); + assert_eq!( + context.tenant_id, + Some("tenant-with-dashes_and_underscores.123".to_string()) + ); +} + +#[test] +fn test_context_tenant_id_empty_value() { + let config = Arc::new(lambda_runtime::Config::default()); + + // Test with empty tenant ID + let mut headers = HeaderMap::new(); + headers.insert("lambda-runtime-aws-request-id", "test-id".parse().unwrap()); + headers.insert("lambda-runtime-deadline-ms", "123456789".parse().unwrap()); + headers.insert("lambda-runtime-aws-tenant-id", "".parse().unwrap()); + + let context = Context::new("test-id", config, &headers).unwrap(); + assert_eq!(context.tenant_id, Some("".to_string())); +} + +#[test] +fn test_trace_layer_request_span_creation() { + use lambda_runtime::layers::trace::request_span; + + // Test with both trace ID and tenant ID + let mut context = Context::default(); + context.request_id = "test-request".to_string(); + context.xray_trace_id = Some("trace-123".to_string()); + context.tenant_id = Some("tenant-456".to_string()); + + let _span = request_span(&context); + // Just verify the span can be created without panicking + + // Test with only trace ID + let mut context = Context::default(); + context.request_id = "test-request".to_string(); + context.xray_trace_id = Some("trace-123".to_string()); + + let _span = request_span(&context); + + // Test with only tenant ID + let mut context = Context::default(); + context.request_id = "test-request".to_string(); + context.tenant_id = Some("tenant-only".to_string()); + + let _span = request_span(&context); + + // Test with only request ID + let mut context = Context::default(); + context.request_id = "test-request".to_string(); + + let _span = request_span(&context); +}