diff --git a/src/main/java/software/amazon/cloudformation/proxy/AmazonWebServicesClientProxy.java b/src/main/java/software/amazon/cloudformation/proxy/AmazonWebServicesClientProxy.java index 76fdef3a..3faf10c6 100644 --- a/src/main/java/software/amazon/cloudformation/proxy/AmazonWebServicesClientProxy.java +++ b/src/main/java/software/amazon/cloudformation/proxy/AmazonWebServicesClientProxy.java @@ -466,7 +466,9 @@ public final long getRemainingTimeInMillis() { request.setRequestCredentialsProvider(v1CredentialsProvider); try { - return requestFunction.apply(request); + ResultT respose = requestFunction.apply(request); + logRequestMetadata(request, respose); + return respose; } catch (final Throwable e) { loggerProxy.log(String.format("Failed to execute remote function: {%s}", e.getMessage())); throw e; @@ -486,7 +488,9 @@ public final long getRemainingTimeInMillis() { RequestT wrappedRequest = (RequestT) request.toBuilder().overrideConfiguration(overrideConfiguration).build(); try { - return requestFunction.apply(wrappedRequest); + ResultT response = requestFunction.apply(wrappedRequest); + logRequestMetadataV2(request, response); + return response; } catch (final Throwable e) { loggerProxy.log(String.format("Failed to execute remote function: {%s}", e.getMessage())); throw e; @@ -505,7 +509,11 @@ public final long getRemainingTimeInMillis() { RequestT wrappedRequest = (RequestT) request.toBuilder().overrideConfiguration(overrideConfiguration).build(); try { - return requestFunction.apply(wrappedRequest); + CompletableFuture response = requestFunction.apply(wrappedRequest).thenApplyAsync(resultT -> { + logRequestMetadataV2(request, resultT); + return resultT; + }); + return response; } catch (final Throwable e) { loggerProxy.log(String.format("Failed to execute remote function: {%s}", e.getMessage())); throw e; @@ -523,7 +531,9 @@ public final long getRemainingTimeInMillis() { RequestT wrappedRequest = (RequestT) request.toBuilder().overrideConfiguration(overrideConfiguration).build(); try { - return requestFunction.apply(wrappedRequest); + IterableT response = requestFunction.apply(wrappedRequest); + response.forEach(r -> logRequestMetadataV2(request, r)); + return response; } catch (final Throwable e) { loggerProxy.log(String.format("Failed to execute remote function: {%s}", e.getMessage())); throw e; @@ -542,7 +552,9 @@ public final long getRemainingTimeInMillis() { RequestT wrappedRequest = (RequestT) request.toBuilder().overrideConfiguration(overrideConfiguration).build(); try { - return requestFunction.apply(wrappedRequest); + ResponseInputStream response = requestFunction.apply(wrappedRequest); + logRequestMetadataV2(request, response.response()); + return response; } catch (final Throwable e) { loggerProxy.log(String.format("Failed to execute remote function: {%s}", e.getMessage())); throw e; @@ -561,7 +573,9 @@ public final long getRemainingTimeInMillis() { RequestT wrappedRequest = (RequestT) request.toBuilder().overrideConfiguration(overrideConfiguration).build(); try { - return requestFunction.apply(wrappedRequest); + ResponseBytes response = requestFunction.apply(wrappedRequest); + logRequestMetadataV2(request, response.response()); + return response; } catch (final Throwable e) { loggerProxy.log(String.format("Failed to execute remote function: {%s}", e.getMessage())); throw e; @@ -623,4 +637,33 @@ public final long getRemainingTimeInMillis() { return ProgressEvent.failed(model, context, HandlerErrorCode.InternalFailure, e.getMessage()); } + + private > + void + logRequestMetadata(final RequestT request, final ResultT response) { + try { + String requestName = request.getClass().getSimpleName(); + String requestId = (response == null || response.getSdkResponseMetadata() == null) + ? "" + : response.getSdkResponseMetadata().getRequestId(); + loggerProxy + .log(String.format("{\"apiRequest\": {\"requestId\": \"%s\", \"requestName\": \"%s\"}}", requestId, requestName)); + } catch (final Exception e) { + loggerProxy.log(e.getMessage()); + } + } + + private void logRequestMetadataV2(final RequestT request, + final ResultT response) { + try { + String requestName = request.getClass().getSimpleName(); + String requestId = (response == null || response.responseMetadata() == null) + ? "" + : response.responseMetadata().requestId(); + loggerProxy + .log(String.format("{\"apiRequest\": {\"requestId\": \"%s\", \"requestName\": \"%s\"}}", requestId, requestName)); + } catch (final Exception e) { + loggerProxy.log(e.getMessage()); + } + } } diff --git a/src/test/java/software/amazon/cloudformation/proxy/AmazonWebServicesClientProxyTest.java b/src/test/java/software/amazon/cloudformation/proxy/AmazonWebServicesClientProxyTest.java index d3dc5181..b4482072 100644 --- a/src/test/java/software/amazon/cloudformation/proxy/AmazonWebServicesClientProxyTest.java +++ b/src/test/java/software/amazon/cloudformation/proxy/AmazonWebServicesClientProxyTest.java @@ -42,17 +42,20 @@ import software.amazon.awssdk.auth.credentials.AwsSessionCredentials; import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider; import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration; +import software.amazon.awssdk.awscore.AwsResponse; import software.amazon.awssdk.awscore.exception.AwsErrorDetails; import software.amazon.awssdk.awscore.exception.AwsServiceException; import software.amazon.awssdk.core.ResponseBytes; import software.amazon.awssdk.core.ResponseInputStream; import software.amazon.awssdk.core.exception.NonRetryableException; +import software.amazon.awssdk.core.pagination.sync.SdkIterable; import software.amazon.awssdk.http.SdkHttpResponse; import software.amazon.awssdk.services.cloudformation.CloudFormationAsyncClient; import software.amazon.awssdk.services.cloudformation.CloudFormationClient; import software.amazon.awssdk.services.cloudformation.model.DescribeStackEventsResponse; import software.amazon.awssdk.services.s3.S3Client; import software.amazon.awssdk.services.s3.model.GetObjectResponse; +import software.amazon.awssdk.services.s3.paginators.ListObjectsV2Iterable; import software.amazon.cloudformation.exceptions.ResourceAlreadyExistsException; import software.amazon.cloudformation.exceptions.TerminalException; import software.amazon.cloudformation.proxy.delay.Constant; @@ -160,6 +163,43 @@ public void testInjectCredentialsAndInvokeV2() { assertThat(result).isEqualTo(expectedResult); } + @Test + public > void testInjectCredentialsAndInvokeV2Iterable() { + + final LoggerProxy loggerProxy = mock(LoggerProxy.class); + final Credentials credentials = new Credentials("accessKeyId", "secretAccessKey", "sessionToken"); + final ListObjectsV2Iterable response = mock(ListObjectsV2Iterable.class); + + final AmazonWebServicesClientProxy proxy = new AmazonWebServicesClientProxy(loggerProxy, credentials, () -> 1000L); + + final software.amazon.awssdk.services.s3.model.ListObjectsV2Request wrappedRequest = mock( + software.amazon.awssdk.services.s3.model.ListObjectsV2Request.class); + + final software.amazon.awssdk.services.s3.model.ListObjectsV2Request.Builder builder = mock( + software.amazon.awssdk.services.s3.model.ListObjectsV2Request.Builder.class); + when(builder.overrideConfiguration(any(AwsRequestOverrideConfiguration.class))).thenReturn(builder); + when(builder.build()).thenReturn(wrappedRequest); + final software.amazon.awssdk.services.s3.model.ListObjectsV2Request request = mock( + software.amazon.awssdk.services.s3.model.ListObjectsV2Request.class); + when(request.toBuilder()).thenReturn(builder); + + final S3Client client = mock(S3Client.class); + + when(client.listObjectsV2Paginator(any(software.amazon.awssdk.services.s3.model.ListObjectsV2Request.class))) + .thenReturn(response); + + final ListObjectsV2Iterable result = proxy.injectCredentialsAndInvokeIterableV2(request, client::listObjectsV2Paginator); + + // verify request is rebuilt for injection + verify(request).toBuilder(); + + // verify the wrapped request is sent over the initiate + verify(client).listObjectsV2Paginator(wrappedRequest); + + // ensure the return type matches + assertThat(result).isEqualTo(response); + } + @Test public void testInjectCredentialsAndInvokeV2Async() throws ExecutionException, InterruptedException {