Java Code Examples for software.amazon.awssdk.core.interceptor.Context#AfterExecution

@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 void afterExecution(Context.AfterExecution context, ExecutionAttributes executionAttributes) {
    try {
        String uri = context.httpRequest().getUri().toString();
        EntityId nodeAccountId = getNodeAccountId(uri);
        Instant startTime = executionAttributes.getAttribute(START_TIME);
        ResponseSizeSubscriber responseSizeSubscriber = executionAttributes.getAttribute(SIZE);

        String[] tags = {
                "action", getAction(uri),
                "method", context.httpRequest().method().name(),
                "nodeAccount", nodeAccountId.getEntityNum().toString(),
                "realm", nodeAccountId.getRealmNum().toString(),
                "shard", nodeAccountId.getShardNum().toString(),
                "status", String.valueOf(context.httpResponse().statusCode()),
                "type", getType(uri)
        };

        if (startTime != null) {
            requestMetric.tags(tags)
                    .register(meterRegistry)
                    .record(Duration.between(startTime, Instant.now()));
        }

        if (responseSizeSubscriber != null) {
            responseSizeMetric.tags(tags)
                    .register(meterRegistry)
                    .record(Double.valueOf(responseSizeSubscriber.getSize()));
        }
    } catch (Exception e) {
        log.warn("Unable to collect S3 metrics", e);
    }
}
 

@Override
public void afterExecution(Context.AfterExecution context, ExecutionAttributes executionAttributes) {
    Subsegment subsegment = executionAttributes.getAttribute(entityKey);
    if (subsegment == null) {
        return;
    }

    populateRequestId(subsegment, context);
    populateSubsegmentWithResponse(subsegment, context.httpResponse());
    subsegment.putAllAws(extractResponseParameters(context, executionAttributes));

    getRecorder().endSubsegment(subsegment);
}
 

/**
 * After a SDK request has been executed
 */
@Override public void afterExecution(
    Context.AfterExecution context,
    ExecutionAttributes executionAttributes
) {
  Span span = executionAttributes.getAttribute(SPAN);
  if (span == null) {
    // An evil interceptor deleted our attribute.
    return;
  }
  handler.handleReceive(
      new HttpClientResponse(context.httpRequest(), context.httpResponse(), null), span);
}
 

private void populateRequestId(Subsegment subsegment, Context.AfterExecution context) {
    populateRequestId(subsegment, Optional.of(context.response()), Optional.of(context.httpResponse()), null);
}
 

private void expectAllMethodsCalled(ExecutionInterceptor interceptor,
                                    SdkRequest inputRequest,
                                    Exception expectedException,
                                    boolean isAsync) {
    ArgumentCaptor<ExecutionAttributes> attributes = ArgumentCaptor.forClass(ExecutionAttributes.class);

    ArgumentCaptor<Context.BeforeExecution> beforeExecutionArg = ArgumentCaptor.forClass(Context.BeforeExecution.class);
    ArgumentCaptor<Context.BeforeMarshalling> modifyRequestArg = ArgumentCaptor.forClass(Context.BeforeMarshalling.class);
    ArgumentCaptor<Context.BeforeMarshalling> beforeMarshallingArg = ArgumentCaptor.forClass(Context.BeforeMarshalling.class);
    ArgumentCaptor<Context.AfterMarshalling> afterMarshallingArg = ArgumentCaptor.forClass(Context.AfterMarshalling.class);
    ArgumentCaptor<Context.BeforeTransmission> modifyHttpRequestArg = ArgumentCaptor.forClass(Context.BeforeTransmission.class);
    ArgumentCaptor<Context.BeforeTransmission> modifyHttpContentArg = ArgumentCaptor.forClass(Context.BeforeTransmission.class);
    ArgumentCaptor<Context.BeforeTransmission> modifyHttpContentAsyncArg = ArgumentCaptor.forClass(Context.BeforeTransmission.class);
    ArgumentCaptor<Context.BeforeTransmission> beforeTransmissionArg = ArgumentCaptor.forClass(Context.BeforeTransmission.class);
    ArgumentCaptor<Context.AfterTransmission> afterTransmissionArg = ArgumentCaptor.forClass(Context.AfterTransmission.class);
    ArgumentCaptor<Context.BeforeUnmarshalling> modifyHttpResponseArg = ArgumentCaptor.forClass(Context.BeforeUnmarshalling.class);
    ArgumentCaptor<Context.BeforeUnmarshalling> modifyHttpResponseContentArg = ArgumentCaptor.forClass(Context.BeforeUnmarshalling.class);
    ArgumentCaptor<Context.BeforeUnmarshalling> beforeUnmarshallingArg = ArgumentCaptor.forClass(Context.BeforeUnmarshalling.class);
    ArgumentCaptor<Context.BeforeUnmarshalling> modifyAsyncHttpResponseContent = ArgumentCaptor.forClass(Context.BeforeUnmarshalling.class);
    ArgumentCaptor<Context.AfterUnmarshalling> afterUnmarshallingArg = ArgumentCaptor.forClass(Context.AfterUnmarshalling.class);
    ArgumentCaptor<Context.AfterExecution> modifyResponseArg = ArgumentCaptor.forClass(Context.AfterExecution.class);
    ArgumentCaptor<Context.AfterExecution> afterExecutionArg = ArgumentCaptor.forClass(Context.AfterExecution.class);

    // Verify methods are called in the right order
    InOrder inOrder = Mockito.inOrder(interceptor);
    inOrder.verify(interceptor).beforeExecution(beforeExecutionArg.capture(), attributes.capture());
    inOrder.verify(interceptor).modifyRequest(modifyRequestArg.capture(), attributes.capture());
    inOrder.verify(interceptor).beforeMarshalling(beforeMarshallingArg.capture(), attributes.capture());
    inOrder.verify(interceptor).afterMarshalling(afterMarshallingArg.capture(), attributes.capture());
    inOrder.verify(interceptor).modifyAsyncHttpContent(modifyHttpContentAsyncArg.capture(), attributes.capture());
    inOrder.verify(interceptor).modifyHttpContent(modifyHttpContentArg.capture(), attributes.capture());
    inOrder.verify(interceptor).modifyHttpRequest(modifyHttpRequestArg.capture(), attributes.capture());
    inOrder.verify(interceptor).beforeTransmission(beforeTransmissionArg.capture(), attributes.capture());
    inOrder.verify(interceptor).afterTransmission(afterTransmissionArg.capture(), attributes.capture());
    inOrder.verify(interceptor).modifyHttpResponse(modifyHttpResponseArg.capture(), attributes.capture());
    inOrder.verify(interceptor).modifyHttpResponseContent(modifyHttpResponseContentArg.capture(), attributes.capture());
    inOrder.verify(interceptor).beforeUnmarshalling(beforeUnmarshallingArg.capture(), attributes.capture());
    if (isAsync) {
        inOrder.verify(interceptor).modifyAsyncHttpResponseContent(modifyAsyncHttpResponseContent.capture(), attributes.capture());
    }
    inOrder.verify(interceptor).afterUnmarshalling(afterUnmarshallingArg.capture(), attributes.capture());
    inOrder.verify(interceptor).modifyResponse(modifyResponseArg.capture(), attributes.capture());
    inOrder.verify(interceptor).afterExecution(afterExecutionArg.capture(), attributes.capture());
    if (expectedException != null) {
        ArgumentCaptor<Context.FailedExecution> failedExecutionArg = ArgumentCaptor.forClass(Context.FailedExecution.class);
        inOrder.verify(interceptor).modifyException(failedExecutionArg.capture(), attributes.capture());
        inOrder.verify(interceptor).onExecutionFailure(failedExecutionArg.capture(), attributes.capture());
        verifyFailedExecutionMethodCalled(failedExecutionArg, true);
        assertThat(failedExecutionArg.getValue().exception()).isEqualTo(expectedException);
    }
    verifyNoMoreInteractions(interceptor);

    // Verify beforeExecution gets untouched request
    assertThat(beforeExecutionArg.getValue().request()).isSameAs(inputRequest);

    // Verify methods were given correct parameters
    validateArgs(beforeExecutionArg.getValue(), null);
    validateArgs(modifyRequestArg.getValue(), null);
    validateArgs(beforeMarshallingArg.getValue(), "1");
    validateArgs(afterMarshallingArg.getValue(), "1", null);
    validateArgs(modifyHttpRequestArg.getValue(), "1", null);
    validateArgs(beforeTransmissionArg.getValue(), "1", "2");
    validateArgs(afterTransmissionArg.getValue(), "1", "2", null);
    validateArgs(modifyHttpResponseArg.getValue(), "1", "2", null);
    validateArgs(beforeUnmarshallingArg.getValue(), "1", "2", "3");
    validateArgs(afterUnmarshallingArg.getValue(), "1", "2", "3", null);
    validateArgs(modifyResponseArg.getValue(), "1", "2", "3", null);
    validateArgs(afterExecutionArg.getValue(), "1", "2", "3", "4");

    // Verify same execution attributes were used for all method calls
    assertThat(attributes.getAllValues()).containsOnly(attributes.getAllValues().get(0));
}
 

@Override
public void afterExecution(Context.AfterExecution context, ExecutionAttributes attrs) {
  attrs.getAttribute(LOG_ENTRY).log();
}