Add multi-tenancy support

.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()

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"] }

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<Value>) -> Result<Value, Error> {
+    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.

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<Value>) -> Result<Value, Error> {
+    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
+}

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"

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<Value>) -> Result<Value, Error> {
+    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
+}

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/"
+    Value: !Sub "https://${API}.execute-api.${AWS::Region}.amazonaws.com/integ-test/hello/"
+  TenantIdTestFunction:
+    Description: "Tenant ID test function name"
+    Value: !Ref TenantIdTestFunction

lambda-integration-tests/tests/tenant_id_prod_test.rs

@@ -0,0 +1,57 @@
+use serde_json::json;
+
+#[test]
+fn test_tenant_id_functionality_in_production() {
+    let function_name =
+        std::env::var("TENANT_ID_TEST_FUNCTION").expect("TENANT_ID_TEST_FUNCTION environment variable not set");
+
+    // Test with tenant ID
+    let payload_with_tenant = json!({
+        "test": "tenant_id_test",
+        "message": "Testing with tenant ID"
+    });
+
+    let output = std::process::Command::new("aws")
+        .args([
+            "lambda",
+            "invoke",
+            "--function-name",
+            &function_name,
+            "--payload",
+            &payload_with_tenant.to_string(),
+            "--cli-binary-format",
+            "raw-in-base64-out",
+            "/tmp/tenant_response.json",
+        ])
+        .output()
+        .expect("Failed to invoke Lambda function");
+
+    assert!(
+        output.status.success(),
+        "Lambda invocation failed: {}",
+        String::from_utf8_lossy(&output.stderr)
+    );
+
+    // Read and verify response
+    let response = std::fs::read_to_string("/tmp/tenant_response.json").expect("Failed to read response file");
+
+    let response_json: serde_json::Value = serde_json::from_str(&response).expect("Failed to parse response JSON");
+
+    // Verify the function executed successfully
+    assert_eq!(response_json["statusCode"], 200);
+
+    // Parse the body to check tenant_id field exists (even if null)
+    let body: serde_json::Value =
+        serde_json::from_str(response_json["body"].as_str().expect("Body should be a string"))
+            .expect("Failed to parse body JSON");
+
+    assert!(
+        body.get("tenant_id").is_some(),
+        "tenant_id field should be present in response"
+    );
+    assert!(body.get("request_id").is_some(), "request_id should be present");
+    assert_eq!(body["message"], "Tenant ID test successful");
+
+    println!("✅ Tenant ID production test passed");
+    println!("Response: {}", serde_json::to_string_pretty(&response_json).unwrap());
+}

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;

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;

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)
         }
     }

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<CognitoIdentity>,
+    /// The tenant ID for the current invocation.
+    pub tenant_id: Option<String>,
     /// 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);
+    }
 }