software.amazon.awssdk.core.SdkResponse Java Examples

/**
 * Combines and decorates separate success and failure response handlers into a single combined response handler
 * that handles both cases and produces a {@link Response} object that wraps the result. The handlers are
 * decorated with additional behavior (such as CRC32 validation).
 */
private <OutputT extends SdkResponse> TransformingAsyncResponseHandler<Response<OutputT>> createDecoratedHandler(
    HttpResponseHandler<OutputT> successHandler,
    HttpResponseHandler<? extends SdkException> errorHandler,
    ExecutionContext executionContext) {

    HttpResponseHandler<OutputT> decoratedResponseHandlers =
        decorateResponseHandlers(successHandler, executionContext);

    TransformingAsyncResponseHandler<OutputT> decoratedSuccessHandler =
        new AsyncResponseHandler<>(decoratedResponseHandlers,
                                   crc32Validator,
                                   executionContext.executionAttributes());

    TransformingAsyncResponseHandler<? extends SdkException> decoratedErrorHandler =
        resolveErrorResponseHandler(errorHandler, executionContext, crc32Validator);
    return new CombinedResponseAsyncHttpResponseHandler<>(decoratedSuccessHandler, decoratedErrorHandler);
}
 

@Test
public void successfulExecutionCallsResponseHandler() throws Exception {

    SdkResponse expected = VoidSdkResponse.builder().build();
    Map<String, List<String>> headers = new HashMap<>();
    headers.put("foo", Arrays.asList("bar"));

    // Given
    expectRetrievalFromMocks();
    when(httpClientCall.call()).thenReturn(HttpExecuteResponse.builder()
                                                              .response(SdkHttpResponse.builder()
                                                                                   .statusCode(200)
                                                                                   .headers(headers)
                                                                                   .build())
                                                              .build()); // Successful HTTP call
    when(responseHandler.handle(any(), any())).thenReturn(expected); // Response handler call

    // When
    SdkResponse actual = syncClientHandler.execute(clientExecutionParams());

    // Then
    verifyNoMoreInteractions(errorResponseHandler); // No error handler calls
    assertThat(actual.sdkHttpResponse().statusCode()).isEqualTo(200);
    assertThat(actual.sdkHttpResponse().headers()).isEqualTo(headers);
}
 

@Test
public void failedExecutionCallsErrorResponseHandler() throws Exception {
    SdkServiceException exception = SdkServiceException.builder().message("Uh oh!").statusCode(500).build();

    // Given
    ArgumentCaptor<AsyncExecuteRequest> executeRequest = ArgumentCaptor.forClass(AsyncExecuteRequest.class);

    expectRetrievalFromMocks();
    when(httpClient.execute(executeRequest.capture())).thenReturn(httpClientFuture);
    when(errorResponseHandler.handle(any(), any())).thenReturn(exception); // Error response handler call

    // When
    CompletableFuture<SdkResponse> responseFuture = asyncClientHandler.execute(clientExecutionParams());
    SdkAsyncHttpResponseHandler capturedHandler = executeRequest.getValue().responseHandler();
    capturedHandler.onHeaders(SdkHttpFullResponse.builder().statusCode(500).build());
    capturedHandler.onStream(new EmptyPublisher<>());
    assertThatThrownBy(() -> responseFuture.get(1, TimeUnit.SECONDS)).hasCause(exception);

    // Then
    verifyNoMoreInteractions(responseHandler); // Response handler is not called
}
 

@Test
public void prepareReturnsNewFuture() {
    AsyncResponseTransformer<SdkResponse, Void> transformer =
            EventStreamAsyncResponseTransformer.builder()
                    .eventStreamResponseHandler(
                            onEventStream(p -> {}))
                    .eventResponseHandler((r, e) -> null)
                    .executor(Executors.newFixedThreadPool(2))
                    .future(new CompletableFuture<>())
                    .build();

    CompletableFuture<?> cf1 = transformer.prepare();

    transformer.exceptionOccurred(new RuntimeException("Boom!"));

    assertThat(cf1.isCompletedExceptionally()).isTrue();
    assertThat(transformer.prepare()).isNotEqualTo(cf1);
}
 

@Test
public void prepareCallsEqualToExecuteAttempts() {
    mockSuccessfulResponse();
    AtomicLong prepareCalls = new AtomicLong(0);
    executeAndWaitError(new TestTransformer<SdkResponse, Void>() {
        @Override
        public CompletableFuture<Void> prepare() {
            prepareCalls.incrementAndGet();
            return super.prepare();
        }

        @Override
        public void onStream(SdkPublisher<ByteBuffer> stream) {
            stream.subscribe(new DrainingSubscriber<ByteBuffer>() {
                @Override
                public void onComplete() {
                    transformFuture().completeExceptionally(RetryableException.builder().message("retry me please: " + prepareCalls.get()).build());
                }
            });
        }
    });
    assertThat(prepareCalls.get()).isEqualTo(1 + RETRY_POLICY.numRetries());
}
 

/**
 * S3 currently returns content-range in two possible headers: Content-Range or x-amz-content-range based on the x-amz-te
 * in the request. This will check the x-amz-content-range if the modeled header (Content-Range) wasn't populated.
 */
private SdkResponse fixContentRange(SdkResponse sdkResponse, SdkHttpResponse httpResponse) {
    // Use the modeled content range header, if the service returned it.
    GetObjectResponse getObjectResponse = (GetObjectResponse) sdkResponse;
    if (getObjectResponse.contentRange() != null) {
        return getObjectResponse;
    }

    // If the service didn't use the modeled content range header, check the x-amz-content-range header.
    Optional<String> xAmzContentRange = httpResponse.firstMatchingHeader("x-amz-content-range");
    if (!xAmzContentRange.isPresent()) {
        return getObjectResponse;
    }

    return getObjectResponse.copy(r -> r.contentRange(xAmzContentRange.get()));
}
 

/**
 * Subtract the contentLength of {@link GetObjectResponse} if trailing checksums is enabled.
 */
@Override
public SdkResponse modifyResponse(Context.ModifyResponse context, ExecutionAttributes executionAttributes) {
    SdkResponse response = context.response();
    SdkHttpResponse httpResponse = context.httpResponse();

    if (getObjectChecksumEnabledPerResponse(context.request(), httpResponse)) {
        GetObjectResponse getResponse = (GetObjectResponse) response;
        Long contentLength = getResponse.contentLength();
        Validate.notNull(contentLength, "Service returned null 'Content-Length'.");
        return getResponse.toBuilder()
                          .contentLength(contentLength - S3_MD5_CHECKSUM_LENGTH)
                          .build();
    }

    return response;
}
 

@Before
public void testSetup() throws Exception {
    mockClient = mock(SdkAsyncHttpClient.class);

    executionParams = new ClientExecutionParams<SdkRequest, SdkResponse>()
            .withInput(request)
            .withMarshaller(marshaller)
            .withResponseHandler(responseHandler)
            .withErrorResponseHandler(errorResponseHandler);

    SdkClientConfiguration config = HttpTestUtils.testClientConfiguration().toBuilder()
            .option(SdkClientOption.ASYNC_HTTP_CLIENT, mockClient)
            .option(SdkClientOption.RETRY_POLICY, RETRY_POLICY)
            .option(SdkAdvancedAsyncClientOption.FUTURE_COMPLETION_EXECUTOR, Runnable::run)
            .build();

    clientHandler = new SdkAsyncClientHandler(config);

    when(request.overrideConfiguration()).thenReturn(Optional.empty());

    when(marshaller.marshall(eq(request))).thenReturn(ValidSdkObjects.sdkHttpFullRequest().build());

    when(responseHandler.handle(any(SdkHttpFullResponse.class), any(ExecutionAttributes.class)))
            .thenReturn(VoidSdkResponse.builder().build());
}
 

/**
 * Run afterUnmarshalling and modifyResponse interceptors.
 */
private static <OutputT extends SdkResponse> BiFunction<OutputT, SdkHttpFullResponse, OutputT>
    runAfterUnmarshallingInterceptors(ExecutionContext context) {

    return (input, httpFullResponse) -> {
        // Update interceptor context to include response
        InterceptorContext interceptorContext =
            context.interceptorContext().copy(b -> b.response(input));

        context.interceptorChain().afterUnmarshalling(interceptorContext, context.executionAttributes());

        interceptorContext = context.interceptorChain().modifyResponse(interceptorContext, context.executionAttributes());

        // Store updated context
        context.interceptorContext(interceptorContext);

        return (OutputT) interceptorContext.response();
    };
}
 

@Override
public <InputT extends SdkRequest, OutputT extends SdkResponse> OutputT execute(
    ClientExecutionParams<InputT, OutputT> executionParams) {

    validateExecutionParams(executionParams);
    ExecutionContext executionContext = createExecutionContext(executionParams, createInitialExecutionAttributes());
    HttpResponseHandler<Response<OutputT>> combinedResponseHandler;

    if (executionParams.getCombinedResponseHandler() != null) {
        combinedResponseHandler = decorateSuccessResponseHandlers(executionParams.getCombinedResponseHandler(),
                                                                  executionContext);
    } else {
        HttpResponseHandler<OutputT> decoratedResponseHandlers =
            decorateResponseHandlers(executionParams.getResponseHandler(), executionContext);

        combinedResponseHandler = new CombinedResponseHandler<>(decoratedResponseHandlers,
                                                                executionParams.getErrorResponseHandler());
    }

    return doExecute(executionParams, executionContext, combinedResponseHandler);
}
 

@Override
public SdkResponse modifyResponse(Context.ModifyResponse context,
                                  ExecutionAttributes executionAttributes) {
    SdkResponse response = context.response();
    if (shouldHandle(response)) {
        if (response instanceof ListObjectsResponse) {
            return modifyListObjectsResponse((ListObjectsResponse) response);
        }

        if (response instanceof ListObjectsV2Response) {
            return modifyListObjectsV2Response((ListObjectsV2Response) response);
        }

        if (response instanceof ListObjectVersionsResponse) {
            return modifyListObjectVersionsResponse((ListObjectVersionsResponse) response);
        }

        if (response instanceof ListMultipartUploadsResponse) {
            return modifyListMultipartUploadsResponse((ListMultipartUploadsResponse) response);
        }
    }
    return response;
}
 

@Override
public void afterExecution(Context.AfterExecution context, ExecutionAttributes executionAttributes) {
    SdkResponse response = context.response();
    SdkRequest originalRequest = context.request();
    if (response != null) {
        if (originalRequest instanceof SendMessageRequest) {
            SendMessageRequest sendMessageRequest = (SendMessageRequest) originalRequest;
            SendMessageResponse sendMessageResult = (SendMessageResponse) response;
            sendMessageOperationMd5Check(sendMessageRequest, sendMessageResult);

        } else if (originalRequest instanceof ReceiveMessageRequest) {
            ReceiveMessageResponse receiveMessageResult = (ReceiveMessageResponse) response;
            receiveMessageResultMd5Check(receiveMessageResult);

        } else if (originalRequest instanceof SendMessageBatchRequest) {
            SendMessageBatchRequest sendMessageBatchRequest = (SendMessageBatchRequest) originalRequest;
            SendMessageBatchResponse sendMessageBatchResult = (SendMessageBatchResponse) response;
            sendMessageBatchOperationMd5Check(sendMessageBatchRequest, sendMessageBatchResult);
        }
    }
}
 

@Override
public <InputT extends SdkRequest, OutputT extends SdkResponse> CompletableFuture<OutputT> execute(
    ClientExecutionParams<InputT, OutputT> executionParams) {

    validateExecutionParams(executionParams);
    ExecutionContext executionContext = createExecutionContext(executionParams, createInitialExecutionAttributes());
    TransformingAsyncResponseHandler<Response<OutputT>> combinedResponseHandler;

    /* Decorate and combine provided response handlers into a single decorated response handler */
    if (executionParams.getCombinedResponseHandler() == null) {
        combinedResponseHandler = createDecoratedHandler(executionParams.getResponseHandler(),
                                                         executionParams.getErrorResponseHandler(),
                                                         executionContext);
    } else {
        combinedResponseHandler = createDecoratedHandler(executionParams.getCombinedResponseHandler(),
                                                         executionContext);
    }

    return doExecute(executionParams, executionContext, combinedResponseHandler);
}
 

private static SdkResponse modifyListObjectsV2Response(ListObjectsV2Response response) {
    return response.toBuilder()
                   .delimiter(urlDecode(response.delimiter()))
                   .prefix(urlDecode(response.prefix()))
                   .startAfter(urlDecode(response.startAfter()))
                   .contents(decodeContents(response.contents()))
                   .commonPrefixes(decodeCommonPrefixes(response.commonPrefixes()))
                   .build();
}
 

@Override
public void onResponse(SdkResponse response) {
    if (response != null && response.sdkHttpResponse() != null) {
        this.requestId = response.sdkHttpResponse()
                                 .firstMatchingHeader(X_AMZN_REQUEST_ID_HEADER)
                                 .orElse(null);
        this.extendedRequestId = response.sdkHttpResponse()
                                         .firstMatchingHeader(X_AMZ_ID_2_HEADER)
                                         .orElse(null);
    }
}
 

private static SdkResponse modifyListObjectsResponse(ListObjectsResponse response) {
    return response.toBuilder()
                   .delimiter(urlDecode(response.delimiter()))
                   .marker(urlDecode(response.marker()))
                   .prefix(urlDecode(response.prefix()))
                   .nextMarker(urlDecode(response.nextMarker()))
                   .contents(decodeContents(response.contents()))
                   .commonPrefixes(decodeCommonPrefixes(response.commonPrefixes()))
                   .build();
}
 

private SdkResponse modifyListObjectVersionsResponse(ListObjectVersionsResponse response) {

        return response.toBuilder()
                       .prefix(urlDecode(response.prefix()))
                       .keyMarker(urlDecode(response.keyMarker()))
                       .delimiter(urlDecode(response.delimiter()))
                       .nextKeyMarker(urlDecode(response.nextKeyMarker()))
                       .commonPrefixes(decodeCommonPrefixes(response.commonPrefixes()))
                       .versions(decodeObjectVersions(response.versions()))
                       .build();
    }
 

/**
 * Add responseHandlers for event streaming operations
 */
private void responseHandlersForEventStreaming(OperationModel opModel, TypeName pojoResponseType,
                                               String protocolFactory, CodeBlock.Builder builder) {
    builder.add("\n\n$T<$T> responseHandler = new $T($L.createResponseHandler(operationMetadata, $T::builder));",
                HttpResponseHandler.class,
                pojoResponseType,
                AttachHttpMetadataResponseHandler.class,
                protocolFactory,
                pojoResponseType);

    builder.add("\n\n$T<$T> voidResponseHandler = $L.createResponseHandler($T.builder()\n" +
                "                                   .isPayloadJson(false)\n" +
                "                                   .hasStreamingSuccessResponse(true)\n" +
                "                                   .build(), $T::builder);",
                HttpResponseHandler.class,
                SdkResponse.class,
                protocolFactory,
                JsonOperationMetadata.class,
                VoidSdkResponse.class);

    ShapeModel eventStream = EventStreamUtils.getEventStreamInResponse(opModel.getOutputShape());
    ClassName eventStreamBaseClass = poetExtensions.getModelClassFromShape(eventStream);
    builder
        .add("\n\n$T<$T> eventResponseHandler = $L.createResponseHandler($T.builder()\n" +
             "                                   .isPayloadJson(true)\n" +
             "                                   .hasStreamingSuccessResponse(false)\n" +
             "                                   .build(), $T.builder()",
             HttpResponseHandler.class,
             WildcardTypeName.subtypeOf(eventStreamBaseClass),
             protocolFactory,
             JsonOperationMetadata.class,
             ClassName.get(EventStreamTaggedUnionPojoSupplier.class));
    EventStreamUtils.getEventMembers(eventStream)
                    .forEach(m -> builder.add(".putSdkPojoSupplier(\"$L\", $T::builder)\n",
                                              m.getC2jName(), poetExtensions.getModelClass(m.getShape().getC2jName())));
    builder.add(".defaultSdkPojoSupplier(() -> new $T($T.UNKNOWN))\n"
                + ".build());\n", SdkPojoBuilder.class, eventStreamBaseClass);
}
 

@Test(timeout = 2000)
public void prepareResetsSubscriberRef() throws InterruptedException {
    CountDownLatch latch = new CountDownLatch(2);
    AtomicBoolean exceptionThrown = new AtomicBoolean(false);

    AsyncResponseTransformer<SdkResponse, Void> transformer =
            EventStreamAsyncResponseTransformer.builder()
                    .eventStreamResponseHandler(
                            onEventStream(p -> {
                                try {
                                    p.subscribe(e -> {});
                                } catch (Throwable t) {
                                    exceptionThrown.set(true);
                                } finally {
                                    latch.countDown();
                                }
                            }))
                    .eventResponseHandler((r, e) -> null)
                    .executor(Executors.newFixedThreadPool(2))
                    .future(new CompletableFuture<>())
                    .build();

    Flowable<ByteBuffer> bytePublisher = Flowable.empty();

    CompletableFuture<Void> transformFuture = transformer.prepare();
    transformer.onStream(SdkPublisher.adapt(bytePublisher));
    transformFuture.join();

    transformFuture = transformer.prepare();
    transformer.onStream(SdkPublisher.adapt(bytePublisher));
    transformFuture.join();

    latch.await();
    assertThat(exceptionThrown).isFalse();
}
 

public UploadFailedException(SdkResponse response) {

        SdkHttpResponse httpResponse = response.sdkHttpResponse();
        if (httpResponse != null) {
            this.statusCode = httpResponse.statusCode();
            this.statusText = httpResponse.statusText();
        } else {
            this.statusCode = HttpStatus.INTERNAL_SERVER_ERROR.value();
            this.statusText = Optional.of("UNKNOWN");
        }

    }
 

public DownloadFailedException(SdkResponse response) {

        SdkHttpResponse httpResponse = response.sdkHttpResponse();
        if (httpResponse != null) {
            this.statusCode = httpResponse.statusCode();
            this.statusText = httpResponse.statusText();
        } else {
            this.statusCode = HttpStatus.INTERNAL_SERVER_ERROR.value();
            this.statusText = Optional.of("UNKNOWN");
        }

    }
 

private void executeAndWaitError(AsyncResponseTransformer<SdkResponse, ?> transformer) {
    try {
        execute(transformer);
        fail("Client execution should have completed exceptionally");
    } catch (CompletionException e) {
        // ignored
    }
}
 

public InterceptorContext modifyResponse(InterceptorContext context, ExecutionAttributes executionAttributes) {
    InterceptorContext result = context;
    for (int i = interceptors.size() - 1; i >= 0; i--) {
        SdkResponse interceptorResult = interceptors.get(i).modifyResponse(result, executionAttributes);
        validateInterceptorResult(result.response(), interceptorResult, interceptors.get(i), "modifyResponse");

        result = result.copy(b -> b.response(interceptorResult));
    }

    return result;
}
 

@Test(timeout = 1000L)
public void handlerExecutionCompletesIndependentlyOfRequestExecution_CompleteAfterStream() {
    mockSuccessfulResponse_NonSignalingStream();
    // Since we never signal any elements on the response stream, if the async client handler waited for the stream to
    // finish before completing the future, this would never return.
    assertThat(execute(new TestTransformer<SdkResponse, Publisher<ByteBuffer>>() {
        @Override
        public void onStream(SdkPublisher<ByteBuffer> stream) {
            transformFuture().complete(stream);
        }
    })).isNotNull();
}
 

/**
 * Returns the composition of 'runAfterUnmarshallingInterceptors' and 'attachHttpResponseToResult' response
 * transformations as a single transformation that should be applied to all responses.
 */
private static <T extends SdkResponse> BiFunction<T, SdkHttpFullResponse, T>
    responseTransformations(ExecutionContext executionContext) {

    return composeResponseFunctions(runAfterUnmarshallingInterceptors(executionContext),
                                    attachHttpResponseToResult());
}
 

@Test
public void marshallerThrowsException_shouldTriggerExceptionOccurred() {
    SdkClientException exception = SdkClientException.create("Could not handle response");
    when(marshaller.marshall(any(SdkRequest.class))).thenThrow(exception);
    AtomicBoolean exceptionOccurred = new AtomicBoolean(false);
    executeAndWaitError(new TestTransformer<SdkResponse, Void>(){
        @Override
        public void exceptionOccurred(Throwable error) {
            exceptionOccurred.set(true);
            super.exceptionOccurred(error);
        }
    });

    assertThat(exceptionOccurred.get()).isTrue();
}
 

<OutputT extends SdkResponse> HttpResponseHandler<OutputT> resultTransformationResponseHandler(
    HttpResponseHandler<OutputT> responseHandler,
    BiFunction<OutputT, SdkHttpFullResponse, OutputT> successTransformer) {

    return (response, executionAttributes) -> {
        OutputT delegateResponse = responseHandler.handle(response, executionAttributes);
        return successTransformer.apply(delegateResponse, response);
    };
}
 

/**
 * Decorates a combined response handler with additional behavior (such as CRC32 validation).
 */
private <OutputT extends SdkResponse> TransformingAsyncResponseHandler<Response<OutputT>> createDecoratedHandler(
    HttpResponseHandler<Response<OutputT>> combinedResponseHandler,
    ExecutionContext executionContext) {

    HttpResponseHandler<Response<OutputT>> decoratedResponseHandlers =
        decorateSuccessResponseHandlers(combinedResponseHandler, executionContext);

    return new AsyncResponseHandler<>(decoratedResponseHandlers,
                                      crc32Validator,
                                      executionContext.executionAttributes());
}
 

private void executeRequestIgnoreErrors(SdkSyncClientHandler syncHttpClient) {
    try {
        syncHttpClient.execute(new ClientExecutionParams<SdkRequest, SdkResponse>()
                .withOperationName("SomeOperation")
                .withResponseHandler(responseHandler)
                .withErrorResponseHandler(errorResponseHandler)
                .withInput(NoopTestRequest.builder().build())
                .withMarshaller(marshaller));
        Assert.fail();
    } catch (AbortedException | ApiCallAttemptTimeoutException | SdkServiceException e) {
        // Ignored
    }
}
 

@Override
public <InputT extends SdkRequest, OutputT extends SdkResponse, ReturnT> ReturnT execute(
    ClientExecutionParams<InputT, OutputT> executionParams,
    ResponseTransformer<OutputT, ReturnT> responseTransformer) {

    validateExecutionParams(executionParams);

    if (executionParams.getCombinedResponseHandler() != null) {
        // There is no support for catching errors in a body for streaming responses
        throw new IllegalArgumentException("A streaming 'responseTransformer' may not be used when a "
                                           + "'combinedResponseHandler' has been specified in a "
                                           + "ClientExecutionParams object.");
    }

    ExecutionContext executionContext = createExecutionContext(executionParams, createInitialExecutionAttributes());

    HttpResponseHandler<OutputT> decoratedResponseHandlers =
        decorateResponseHandlers(executionParams.getResponseHandler(), executionContext);

    HttpResponseHandler<ReturnT> httpResponseHandler =
        new HttpResponseHandlerAdapter<>(decoratedResponseHandlers, responseTransformer);

    return doExecute(
        executionParams,
        executionContext,
        new CombinedResponseHandler<>(httpResponseHandler, executionParams.getErrorResponseHandler()));
}