Skip to content

Commit b7c13ac

Browse files
author
Ramisa Alam
committed
feat: add tenant id to lambda context and structured log message
1 parent c5a4d2e commit b7c13ac

File tree

11 files changed

+127
-13
lines changed

11 files changed

+127
-13
lines changed

Diff for: ‎aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/EventHandlerLoader.java

+1
Original file line numberDiff line numberDiff line change
@@ -581,6 +581,7 @@ public ByteArrayOutputStream call(InvocationRequest request) throws Error, Excep
581581
cognitoIdentity,
582582
LambdaEnvironment.FUNCTION_VERSION,
583583
request.getInvokedFunctionArn(),
584+
request.getTenantId(),
584585
clientContext
585586
);
586587

Diff for: ‎aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/api/LambdaContext.java

+7
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ public class LambdaContext implements Context {
2222
private final long deadlineTimeInMs;
2323
private final CognitoIdentity cognitoIdentity;
2424
private final ClientContext clientContext;
25+
private final String tenantId;
2526
private final LambdaLogger logger;
2627

2728
public LambdaContext(
@@ -34,6 +35,7 @@ public LambdaContext(
3435
CognitoIdentity identity,
3536
String functionVersion,
3637
String invokedFunctionArn,
38+
String tenantId,
3739
ClientContext clientContext
3840
) {
3941
this.memoryLimit = memoryLimit;
@@ -46,6 +48,7 @@ public LambdaContext(
4648
this.clientContext = clientContext;
4749
this.functionVersion = functionVersion;
4850
this.invokedFunctionArn = invokedFunctionArn;
51+
this.tenantId = tenantId;
4952
this.logger = com.amazonaws.services.lambda.runtime.LambdaRuntime.getLogger();
5053
}
5154

@@ -91,6 +94,10 @@ public int getRemainingTimeInMillis() {
9194
return delta > 0 ? delta : 0;
9295
}
9396

97+
public String getTenantId() {
98+
return tenantId;
99+
}
100+
94101
public LambdaLogger getLogger() {
95102
return logger;
96103
}

Diff for: ‎aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/logging/JsonLogFormatter.java

+1
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ private StructuredLogMessage createLogMessage(String message, LogLevel logLevel)
4141

4242
if (lambdaContext != null) {
4343
msg.AWSRequestId = lambdaContext.getAwsRequestId();
44+
msg.tenantId = lambdaContext.getTenantId();
4445
}
4546
return msg;
4647
}

Diff for: ‎aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/logging/StructuredLogMessage.java

+1
Original file line numberDiff line numberDiff line change
@@ -12,4 +12,5 @@ class StructuredLogMessage {
1212
public String message;
1313
public LogLevel level;
1414
public String AWSRequestId;
15+
public String tenantId;
1516
}

Diff for: ‎aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/runtimeapi/dto/InvocationRequest.java

+13
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,11 @@ public class InvocationRequest {
4040
*/
4141
private String cognitoIdentity;
4242

43+
/**
44+
* The tenant ID associated with the request.
45+
*/
46+
private String tenantId;
47+
4348
private byte[] content;
4449

4550
public String getId() {
@@ -94,6 +99,14 @@ public void setCognitoIdentity(String cognitoIdentity) {
9499
this.cognitoIdentity = cognitoIdentity;
95100
}
96101

102+
public String getTenantId() {
103+
return tenantId;
104+
}
105+
106+
public void setTenantId(String tenantId) {
107+
this.tenantId = tenantId;
108+
}
109+
97110
public byte[] getContent() {
98111
return content;
99112
}

Diff for: ‎aws-lambda-java-runtime-interface-client/src/main/jni/com_amazonaws_services_lambda_runtime_api_client_runtimeapi_NativeClient.cpp

+6
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ static jfieldID contentField;
2020
static jfieldID clientContextField;
2121
static jfieldID cognitoIdentityField;
2222
static jfieldID xrayTraceIdField;
23+
static jfieldID tenantIdField;
2324

2425

2526
jint JNI_OnLoad(JavaVM* vm, void* reserved) {
@@ -41,6 +42,7 @@ jint JNI_OnLoad(JavaVM* vm, void* reserved) {
4142
xrayTraceIdField = env->GetFieldID(invocationRequestClass , "xrayTraceId", "Ljava/lang/String;");
4243
clientContextField = env->GetFieldID(invocationRequestClass , "clientContext", "Ljava/lang/String;");
4344
cognitoIdentityField = env->GetFieldID(invocationRequestClass , "cognitoIdentity", "Ljava/lang/String;");
45+
tenantIdField = env->GetFieldID(invocationRequestClass, "tenantId", "Ljava/lang/String;");
4446

4547
return JNI_VERSION;
4648
}
@@ -106,6 +108,10 @@ JNIEXPORT jobject JNICALL Java_com_amazonaws_services_lambda_runtime_api_client_
106108
CHECK_EXCEPTION(env, env->SetObjectField(invocationRequest, cognitoIdentityField, env->NewStringUTF(response.cognito_identity.c_str())));
107109
}
108110

111+
if(response.tenant_id != ""){
112+
CHECK_EXCEPTION(env, env->SetObjectField(invocationRequest, tenantIdField, env->NewStringUTF(response.tenant_id.c_str())));
113+
}
114+
109115
bytes = reinterpret_cast<const jbyte*>(response.payload.c_str());
110116
CHECK_EXCEPTION(env, jArray = env->NewByteArray(response.payload.length()));
111117
CHECK_EXCEPTION(env, env->SetByteArrayRegion(jArray, 0, response.payload.length(), bytes));

Diff for: ‎aws-lambda-java-runtime-interface-client/src/main/jni/deps/aws-lambda-cpp-0.2.7/include/aws/lambda-runtime/runtime.h

+5
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,11 @@ struct invocation_request {
6161
*/
6262
std::chrono::time_point<std::chrono::system_clock> deadline;
6363

64+
/**
65+
* Tenant ID of the current invocation.
66+
*/
67+
std::string tenant_id;
68+
6469
/**
6570
* The number of milliseconds left before lambda terminates the current execution.
6671
*/

Diff for: ‎aws-lambda-java-runtime-interface-client/src/main/jni/deps/aws-lambda-cpp-0.2.7/src/runtime.cpp

+5
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ static constexpr auto CLIENT_CONTEXT_HEADER = "lambda-runtime-client-context";
4040
static constexpr auto COGNITO_IDENTITY_HEADER = "lambda-runtime-cognito-identity";
4141
static constexpr auto DEADLINE_MS_HEADER = "lambda-runtime-deadline-ms";
4242
static constexpr auto FUNCTION_ARN_HEADER = "lambda-runtime-invoked-function-arn";
43+
static constexpr auto TENANT_ID_HEADER = "lambda-runtime-aws-tenant-id";
4344

4445
enum Endpoints {
4546
INIT,
@@ -301,6 +302,10 @@ runtime::next_outcome runtime::get_next()
301302
req.payload.c_str(),
302303
static_cast<int64_t>(req.get_time_remaining().count()));
303304
}
305+
306+
if (resp.has_header(TENANT_ID_HEADER)) {
307+
req.tenant_id = resp.get_header(TENANT_ID_HEADER);
308+
}
304309
return next_outcome(req);
305310
}
306311

Diff for: ‎aws-lambda-java-runtime-interface-client/src/test/java/com/amazonaws/services/lambda/runtime/api/client/api/LambdaContextTest.java

+2-1
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ public class LambdaContextTest {
1818
private static final String INVOKED_FUNCTION_ARN = "invoked-function-arn";
1919
private static final LambdaClientContext CLIENT_CONTEXT = new LambdaClientContext();
2020
public static final int MEMORY_LIMIT = 128;
21+
public static final String TENANT_ID = "tenant-id";
2122

2223
@Test
2324
public void getRemainingTimeInMillis() {
@@ -54,6 +55,6 @@ public void getRemainingTimeInMillis_Deadline() throws InterruptedException {
5455

5556
private LambdaContext createContextWithDeadline(long deadlineTimeInMs) {
5657
return new LambdaContext(MEMORY_LIMIT, deadlineTimeInMs, REQUEST_ID, LOG_GROUP_NAME, LOG_STREAM_NAME,
57-
FUNCTION_NAME, IDENTITY, FUNCTION_VERSION, INVOKED_FUNCTION_ARN, CLIENT_CONTEXT);
58+
FUNCTION_NAME, IDENTITY, FUNCTION_VERSION, INVOKED_FUNCTION_ARN, TENANT_ID, CLIENT_CONTEXT);
5859
}
5960
}

Diff for: ‎aws-lambda-java-runtime-interface-client/src/test/java/com/amazonaws/services/lambda/runtime/api/client/logging/JsonLogFormatterTest.java

+20
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,25 @@ void testFormattingWithLambdaContext() {
2929
null,
3030
null,
3131
"function-arn",
32+
null,
33+
null
34+
);
35+
assertFormatsString("test log", LogLevel.WARN, context);
36+
}
37+
38+
@Test
39+
void testFormattingWithTenantIdInLambdaContext() {
40+
LambdaContext context = new LambdaContext(
41+
0,
42+
0,
43+
"request-id",
44+
null,
45+
null,
46+
"function-name",
47+
null,
48+
null,
49+
"function-arn",
50+
"tenant-id",
3251
null
3352
);
3453
assertFormatsString("test log", LogLevel.WARN, context);
@@ -52,6 +71,7 @@ void assert_expected_log_message(StructuredLogMessage result, String message, Lo
5271

5372
if (context != null) {
5473
assertEquals(context.getAwsRequestId(), result.AWSRequestId);
74+
assertEquals(context.getTenantId(), result.tenantId);
5575
}
5676
}
5777
}

Diff for: ‎aws-lambda-java-runtime-interface-client/src/test/java/com/amazonaws/services/lambda/runtime/api/client/runtimeapi/LambdaRuntimeApiClientImplTest.java

+66-12
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
import static org.junit.jupiter.api.Assertions.assertEquals;
1616
import static org.junit.jupiter.api.Assertions.assertThrows;
1717
import com.amazonaws.services.lambda.runtime.api.client.runtimeapi.dto.ErrorRequest;
18+
import com.amazonaws.services.lambda.runtime.api.client.runtimeapi.dto.InvocationRequest;
1819
import com.amazonaws.services.lambda.runtime.api.client.runtimeapi.dto.StackElement;
1920
import com.amazonaws.services.lambda.runtime.api.client.runtimeapi.dto.XRayErrorCause;
2021
import com.amazonaws.services.lambda.runtime.api.client.runtimeapi.dto.XRayException;
@@ -312,27 +313,62 @@ public void restoreNextWrongStatusCodeTest() {
312313
}
313314

314315
@Test
315-
public void nextTest() {
316+
public void nextWithoutTenantIdHeaderTest() {
316317
try {
317-
MockResponse mockResponse = new MockResponse();
318-
mockResponse.setResponseCode(HTTP_ACCEPTED);
319-
mockResponse.setHeader("lambda-runtime-aws-request-id", "1234567890");
320-
mockResponse.setHeader("Content-Type", "application/json");
318+
MockResponse mockResponse = buildMockResponseForNextInvocation();
321319
mockWebServer.enqueue(mockResponse);
322320

323-
lambdaRuntimeApiClientImpl.nextInvocation();
324-
RecordedRequest recordedRequest = mockWebServer.takeRequest();
325-
HttpUrl actualUrl = recordedRequest.getRequestUrl();
326-
String expectedUrl = "http://" + getHostnamePort() + "/2018-06-01/runtime/invocation/next";
327-
assertEquals(expectedUrl, actualUrl.toString());
321+
InvocationRequest invocationRequest = lambdaRuntimeApiClientImpl.nextInvocation();
322+
verifyNextInvocationRequest();
323+
assertNull(invocationRequest.getTenantId());
324+
} catch(Exception e) {
325+
fail();
326+
}
327+
}
328+
329+
@Test
330+
public void nextWithTenantIdHeaderTest() {
331+
try {
332+
MockResponse mockResponse = buildMockResponseForNextInvocation();
333+
String expectedTenantId = "my-tenant-id";
334+
mockResponse.setHeader("lambda-runtime-aws-tenant-id", expectedTenantId);
335+
mockWebServer.enqueue(mockResponse);
336+
337+
InvocationRequest invocationRequest = lambdaRuntimeApiClientImpl.nextInvocation();
338+
verifyNextInvocationRequest();
339+
assertEquals(expectedTenantId, invocationRequest.getTenantId());
328340

329-
String actualBody = recordedRequest.getBody().readUtf8();
330-
assertEquals("", actualBody);
331341
} catch(Exception e) {
332342
fail();
333343
}
334344
}
335345

346+
@Test
347+
public void nextWithEmptyTenantIdHeaderTest() {
348+
try {
349+
MockResponse mockResponse = buildMockResponseForNextInvocation();
350+
mockResponse.setHeader("lambda-runtime-aws-tenant-id", "");
351+
mockWebServer.enqueue(mockResponse);
352+
353+
InvocationRequest invocationRequest = lambdaRuntimeApiClientImpl.nextInvocation();
354+
verifyNextInvocationRequest();
355+
assertNull(invocationRequest.getTenantId());
356+
} catch(Exception e) {
357+
fail();
358+
}
359+
}
360+
361+
@Test
362+
public void nextWithNullTenantIdHeaderTest() {
363+
try {
364+
MockResponse mockResponse = buildMockResponseForNextInvocation();
365+
assertThrows(NullPointerException.class, () -> {
366+
mockResponse.setHeader("lambda-runtime-aws-tenant-id", null);
367+
});
368+
} catch(Exception e) {
369+
fail();
370+
}
371+
}
336372

337373
@Test
338374
public void createUrlMalformedTest() {
@@ -376,6 +412,24 @@ public void lambdaReportErrorXRayHeaderTooLongTest() {
376412
}
377413
}
378414

415+
private MockResponse buildMockResponseForNextInvocation() {
416+
MockResponse mockResponse = new MockResponse();
417+
mockResponse.setResponseCode(HTTP_ACCEPTED);
418+
mockResponse.setHeader("lambda-runtime-aws-request-id", "1234567890");
419+
mockResponse.setHeader("Content-Type", "application/json");
420+
return mockResponse;
421+
}
422+
423+
private void verifyNextInvocationRequest() throws Exception {
424+
RecordedRequest recordedRequest = mockWebServer.takeRequest();
425+
HttpUrl actualUrl = recordedRequest.getRequestUrl();
426+
String expectedUrl = "http://" + getHostnamePort() + "/2018-06-01/runtime/invocation/next";
427+
assertEquals(expectedUrl, actualUrl.toString());
428+
429+
String actualBody = recordedRequest.getBody().readUtf8();
430+
assertEquals("", actualBody);
431+
}
432+
379433
private String getHostnamePort() {
380434
return mockWebServer.getHostName() + ":" + mockWebServer.getPort();
381435
}

0 commit comments

Comments
 (0)