GH-913 Fix AWS Context propagation

Resolves #913
This commit is contained in:
Oleg Zhurakousky
2022-08-08 15:12:51 +02:00
parent c5948bd07c
commit fc50ebc870
4 changed files with 131 additions and 8 deletions

View File

@@ -22,6 +22,7 @@ import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicReference;
import com.amazonaws.services.lambda.runtime.Context;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
@@ -65,8 +66,12 @@ final class AWSLambdaUtils {
|| typeName.equals("com.amazonaws.services.lambda.runtime.events.KinesisEvent");
}
@SuppressWarnings({ "unchecked", "rawtypes" })
public static Message<byte[]> generateMessage(byte[] payload, Type inputType, boolean isSupplier, JsonMapper jsonMapper) {
return generateMessage(payload, inputType, isSupplier, jsonMapper, null);
}
@SuppressWarnings({ "unchecked", "rawtypes" })
public static Message<byte[]> generateMessage(byte[] payload, Type inputType, boolean isSupplier, JsonMapper jsonMapper, Context context) {
if (logger.isInfoEnabled()) {
logger.info("Received: " + new String(payload, StandardCharsets.UTF_8));
}
@@ -84,6 +89,9 @@ final class AWSLambdaUtils {
if (!isSupplier && AWSLambdaUtils.isSupportedAWSType(inputType)) {
builder.setHeader(AWSLambdaUtils.AWS_EVENT, true);
}
if (context != null) {
builder.setHeader(AWSLambdaUtils.AWS_CONTEXT, context);
}
//
if (structMessage instanceof Map && ((Map<String, Object>) structMessage).containsKey("headers")) {
builder.copyHeaders((Map<String, Object>) ((Map<String, Object>) structMessage).get("headers"));

View File

@@ -80,7 +80,7 @@ public class FunctionInvoker implements RequestStreamHandler {
@Override
public void handleRequest(InputStream input, OutputStream output, Context context) throws IOException {
Message requestMessage = AWSLambdaUtils
.generateMessage(StreamUtils.copyToByteArray(input), this.function.getInputType(), this.function.isSupplier(), jsonMapper);
.generateMessage(StreamUtils.copyToByteArray(input), this.function.getInputType(), this.function.isSupplier(), jsonMapper, context);
Object response = this.function.apply(requestMessage);
byte[] responseBytes = this.buildResult(requestMessage, response);