Add initial support for unwrapping reactive respponses during AWS processing

This commit is contained in:
Anton Barkan
2023-12-26 22:45:35 +02:00
committed by Oleg Zhurakousky
parent 0f6aaaec45
commit 8901c7b5e1
2 changed files with 43 additions and 11 deletions

View File

@@ -70,12 +70,13 @@ public final class AWSLambdaUtils {
} }
static boolean isSupportedAWSType(Type inputType) { static boolean isSupportedAWSType(Type type) {
if (FunctionTypeUtils.isMessage(inputType) || FunctionTypeUtils.isPublisher(inputType)) { if (FunctionTypeUtils.isMessage(type) || FunctionTypeUtils.isPublisher(type)) {
inputType = FunctionTypeUtils.getImmediateGenericType(inputType, 0); type = FunctionTypeUtils.getImmediateGenericType(type, 0);
} }
return FunctionTypeUtils.getRawType(inputType).getPackage() != null && Class<?> rawType = FunctionTypeUtils.getRawType(type);
FunctionTypeUtils.getRawType(inputType).getPackage().getName().startsWith( return rawType != null && rawType.getPackage() != null &&
rawType.getPackage().getName().startsWith(
"com.amazonaws.services.lambda.runtime.events"); "com.amazonaws.services.lambda.runtime.events");
} }
@@ -203,12 +204,8 @@ public final class AWSLambdaUtils {
public static byte[] generateOutput(Message requestMessage, Message<?> responseMessage, public static byte[] generateOutput(Message requestMessage, Message<?> responseMessage,
JsonMapper objectMapper, Type functionOutputType) { JsonMapper objectMapper, Type functionOutputType) {
Class<?> outputClass = FunctionTypeUtils.getRawType(functionOutputType); if (isSupportedAWSType(functionOutputType)) {
if (outputClass != null) { return extractPayload((Message<Object>) responseMessage, objectMapper);
String outputClassName = outputClass.getName();
if (outputClassName.startsWith("com.amazonaws.services.lambda.runtime.events.")) {
return extractPayload((Message<Object>) responseMessage, objectMapper);
}
} }
byte[] responseBytes = responseMessage == null ? "\"OK\"".getBytes() : extractPayload((Message<Object>) responseMessage, objectMapper); byte[] responseBytes = responseMessage == null ? "\"OK\"".getBytes() : extractPayload((Message<Object>) responseMessage, objectMapper);

View File

@@ -24,6 +24,7 @@ import java.io.OutputStream;
import java.nio.charset.StandardCharsets; import java.nio.charset.StandardCharsets;
import java.util.Base64; import java.util.Base64;
import java.util.Collections; import java.util.Collections;
import java.util.List;
import java.util.Map; import java.util.Map;
import java.util.function.Consumer; import java.util.function.Consumer;
import java.util.function.Function; import java.util.function.Function;
@@ -38,6 +39,7 @@ import com.amazonaws.services.lambda.runtime.events.APIGatewayV2HTTPResponse;
import com.amazonaws.services.lambda.runtime.events.ApplicationLoadBalancerRequestEvent; import com.amazonaws.services.lambda.runtime.events.ApplicationLoadBalancerRequestEvent;
import com.amazonaws.services.lambda.runtime.events.ApplicationLoadBalancerResponseEvent; import com.amazonaws.services.lambda.runtime.events.ApplicationLoadBalancerResponseEvent;
import com.amazonaws.services.lambda.runtime.events.DynamodbEvent; import com.amazonaws.services.lambda.runtime.events.DynamodbEvent;
import com.amazonaws.services.lambda.runtime.events.IamPolicyResponse;
import com.amazonaws.services.lambda.runtime.events.KinesisEvent; import com.amazonaws.services.lambda.runtime.events.KinesisEvent;
import com.amazonaws.services.lambda.runtime.events.S3Event; import com.amazonaws.services.lambda.runtime.events.S3Event;
import com.amazonaws.services.lambda.runtime.events.SNSEvent; import com.amazonaws.services.lambda.runtime.events.SNSEvent;
@@ -1347,6 +1349,21 @@ public class FunctionInvokerTests {
assertThat(result.get("body")).isEqualTo("\"hello\""); assertThat(result.get("body")).isEqualTo("\"hello\"");
} }
@Test
public void testShouldNotWrapIamPolicyResponse() throws Exception {
System.setProperty("MAIN_CLASS", ApiGatewayConfiguration.class.getName());
System.setProperty("spring.cloud.function.definition", "outputPolicyResponse");
FunctionInvoker invoker = new FunctionInvoker();
InputStream targetStream = new ByteArrayInputStream(this.apiGatewayEvent.getBytes());
ByteArrayOutputStream output = new ByteArrayOutputStream();
invoker.handleRequest(targetStream, output, null);
Map result = mapper.readValue(output.toByteArray(), Map.class);
assertThat(result.get("body")).isNull();
assertThat(result.get("principalId")).isNotNull();
}
@SuppressWarnings("rawtypes") @SuppressWarnings("rawtypes")
@Test @Test
public void testApiGatewayEventConsumer() throws Exception { public void testApiGatewayEventConsumer() throws Exception {
@@ -1829,6 +1846,24 @@ public class FunctionInvokerTests {
return body; return body;
}; };
} }
@Bean
public Function<Mono<String>, Mono<IamPolicyResponse>> outputPolicyResponse() {
return input ->
input.map(v -> IamPolicyResponse.builder()
.withPrincipalId("principalId")
.withPolicyDocument(IamPolicyResponse.PolicyDocument.builder()
.withVersion("2012-10-17")
.withStatement(
List.of(
IamPolicyResponse.Statement.builder().withAction("execute-api:Invoke")
.withResource(
List.of(v)).withEffect("Allow").build()
)
).build()
).build()
);
}
} }
@EnableAutoConfiguration @EnableAutoConfiguration