Add multi-tenancy support

Author: PartiallyUntyped

.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,49 @@
+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;