Skip to content

Commit 57a47d5

Browse files
QuanzzzzQuanzzzz
committed
Append userAgent for Waiters operations
1 parent b154ef3 commit 57a47d5

File tree

6 files changed

+189
-5
lines changed

6 files changed

+189
-5
lines changed

‎codegen/src/main/java/software/amazon/awssdk/codegen/poet/waiters/BaseWaiterClassSpec.java

Lines changed: 45 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
import com.squareup.javapoet.ParameterizedTypeName;
3030
import com.squareup.javapoet.TypeName;
3131
import com.squareup.javapoet.TypeSpec;
32+
import com.squareup.javapoet.TypeVariableName;
3233
import com.squareup.javapoet.WildcardTypeName;
3334
import java.time.Duration;
3435
import java.util.ArrayList;
@@ -37,19 +38,23 @@
3738
import java.util.Map;
3839
import java.util.Objects;
3940
import java.util.Optional;
41+
import java.util.function.Consumer;
4042
import java.util.stream.Collectors;
4143
import java.util.stream.Stream;
4244
import javax.lang.model.element.Modifier;
4345
import software.amazon.awssdk.annotations.SdkInternalApi;
4446
import software.amazon.awssdk.annotations.ThreadSafe;
47+
import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
4548
import software.amazon.awssdk.awscore.exception.AwsServiceException;
4649
import software.amazon.awssdk.codegen.emitters.tasks.WaitersRuntimeGeneratorTask;
4750
import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
4851
import software.amazon.awssdk.codegen.model.intermediate.OperationModel;
4952
import software.amazon.awssdk.codegen.model.service.Acceptor;
5053
import software.amazon.awssdk.codegen.model.service.WaiterDefinition;
5154
import software.amazon.awssdk.codegen.poet.ClassSpec;
55+
import software.amazon.awssdk.codegen.poet.PoetExtensions;
5256
import software.amazon.awssdk.codegen.poet.PoetUtils;
57+
import software.amazon.awssdk.core.ApiName;
5358
import software.amazon.awssdk.core.internal.waiters.WaiterAttribute;
5459
import software.amazon.awssdk.core.retry.backoff.FixedDelayBackoffStrategy;
5560
import software.amazon.awssdk.core.waiters.PollingStrategy;
@@ -62,18 +67,22 @@
6267
* Base class containing common logic shared between the sync waiter class and the async waiter class
6368
*/
6469
public abstract class BaseWaiterClassSpec implements ClassSpec {
70+
71+
private static final String WAITERS_USER_AGENT = "waiter";
6572
private final IntermediateModel model;
6673
private final String modelPackage;
6774
private final Map<String, WaiterDefinition> waiters;
6875
private final ClassName waiterClassName;
6976
private final JmesPathAcceptorGenerator jmesPathAcceptorGenerator;
77+
private final PoetExtensions poetExtensions;
7078

7179
public BaseWaiterClassSpec(IntermediateModel model, ClassName waiterClassName) {
7280
this.model = model;
7381
this.modelPackage = model.getMetadata().getFullModelPackageName();
7482
this.waiters = model.getWaiters();
7583
this.waiterClassName = waiterClassName;
7684
this.jmesPathAcceptorGenerator = new JmesPathAcceptorGenerator(waitersRuntimeClass());
85+
this.poetExtensions = new PoetExtensions(model);
7786
}
7887

7988
@Override
@@ -105,6 +114,7 @@ public TypeSpec poetSpec() {
105114
.build());
106115

107116
typeSpecBuilder.addType(builder());
117+
typeSpecBuilder.addMethod(applyWaitersUserAgentMethod(poetExtensions, model));
108118
return typeSpecBuilder.build();
109119
}
110120

@@ -277,7 +287,7 @@ private MethodSpec waiterOperation(Map.Entry<String, WaiterDefinition> waiterDef
277287
.addParameter(requestType, opModel.getInput().getVariableName())
278288
.addModifiers(PUBLIC)
279289
.addAnnotation(Override.class)
280-
.addStatement("return $L.$L(() -> client.$N($N))",
290+
.addStatement("return $L.$L(() -> client.$N(applyWaitersUserAgent($N)))",
281291
waiterFieldName(waiterMethodName),
282292
waiterClassName.simpleName().equals("Waiter") ? "run" : "runAsync",
283293
lowercaseFirstChar(waiterDefinition.getValue().getOperation()),
@@ -326,6 +336,40 @@ private MethodSpec.Builder methodSignatureWithReturnType(String waiterMethodName
326336
.returns(getWaiterResponseType(opModel));
327337
}
328338

339+
static MethodSpec applyWaitersUserAgentMethod(PoetExtensions poetExtensions, IntermediateModel model) {
340+
341+
TypeVariableName typeVariableName =
342+
TypeVariableName.get("T", poetExtensions.getModelClass(model.getSdkRequestBaseClassName()));
343+
344+
ParameterizedTypeName parameterizedTypeName = ParameterizedTypeName
345+
.get(ClassName.get(Consumer.class), ClassName.get(AwsRequestOverrideConfiguration.Builder.class));
346+
347+
CodeBlock codeBlock = CodeBlock.builder()
348+
.addStatement("$T userAgentApplier = b -> b.addApiName($T.builder().version"
349+
+ "($S).name($S).build())",
350+
parameterizedTypeName, ApiName.class,
351+
WAITERS_USER_AGENT,
352+
"hll")
353+
.addStatement("$T overrideConfiguration =\n"
354+
+ " request.overrideConfiguration().map(c -> c.toBuilder()"
355+
+ ".applyMutation"
356+
+ "(userAgentApplier).build())\n"
357+
+ " .orElse((AwsRequestOverrideConfiguration.builder()"
358+
+ ".applyMutation"
359+
+ "(userAgentApplier).build()))", AwsRequestOverrideConfiguration.class)
360+
.addStatement("return (T) request.toBuilder().overrideConfiguration"
361+
+ "(overrideConfiguration).build()")
362+
.build();
363+
364+
return MethodSpec.methodBuilder("applyWaitersUserAgent")
365+
.addModifiers(Modifier.PRIVATE)
366+
.addParameter(typeVariableName, "request")
367+
.addTypeVariable(typeVariableName)
368+
.addCode(codeBlock)
369+
.returns(typeVariableName)
370+
.build();
371+
}
372+
329373
private String getWaiterMethodName(String waiterMethodName) {
330374
return "waitUntil" + waiterMethodName;
331375
}

‎codegen/src/test/resources/software/amazon/awssdk/codegen/poet/waiters/query-async-waiter-class.java

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,13 @@
77
import java.util.concurrent.CompletableFuture;
88
import java.util.concurrent.Executors;
99
import java.util.concurrent.ScheduledExecutorService;
10+
import java.util.function.Consumer;
1011
import software.amazon.awssdk.annotations.Generated;
1112
import software.amazon.awssdk.annotations.SdkInternalApi;
1213
import software.amazon.awssdk.annotations.ThreadSafe;
14+
import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
1315
import software.amazon.awssdk.awscore.exception.AwsServiceException;
16+
import software.amazon.awssdk.core.ApiName;
1417
import software.amazon.awssdk.core.internal.waiters.WaiterAttribute;
1518
import software.amazon.awssdk.core.retry.backoff.FixedDelayBackoffStrategy;
1619
import software.amazon.awssdk.core.waiters.AsyncWaiter;
@@ -21,6 +24,7 @@
2124
import software.amazon.awssdk.services.query.QueryAsyncClient;
2225
import software.amazon.awssdk.services.query.model.APostOperationRequest;
2326
import software.amazon.awssdk.services.query.model.APostOperationResponse;
27+
import software.amazon.awssdk.services.query.model.QueryRequest;
2428
import software.amazon.awssdk.services.query.waiters.internal.WaitersRuntime;
2529
import software.amazon.awssdk.utils.AttributeMap;
2630
import software.amazon.awssdk.utils.SdkAutoCloseable;
@@ -77,7 +81,7 @@ private static String errorCode(Throwable error) {
7781
@Override
7882
public CompletableFuture<WaiterResponse<APostOperationResponse>> waitUntilPostOperationSuccess(
7983
APostOperationRequest aPostOperationRequest) {
80-
return postOperationSuccessWaiter.runAsync(() -> client.aPostOperation(aPostOperationRequest));
84+
return postOperationSuccessWaiter.runAsync(() -> client.aPostOperation(applyWaitersUserAgent(aPostOperationRequest)));
8185
}
8286

8387
private static List<WaiterAcceptor<? super APostOperationResponse>> postOperationSuccessWaiterAcceptors() {
@@ -102,6 +106,15 @@ public static QueryAsyncWaiter.Builder builder() {
102106
return new DefaultBuilder();
103107
}
104108

109+
private <T extends QueryRequest> T applyWaitersUserAgent(T request) {
110+
Consumer<AwsRequestOverrideConfiguration.Builder> userAgentApplier = b -> b.addApiName(ApiName.builder()
111+
.version("waiter").name("hll").build());
112+
AwsRequestOverrideConfiguration overrideConfiguration = request.overrideConfiguration()
113+
.map(c -> c.toBuilder().applyMutation(userAgentApplier).build())
114+
.orElse((AwsRequestOverrideConfiguration.builder().applyMutation(userAgentApplier).build()));
115+
return (T) request.toBuilder().overrideConfiguration(overrideConfiguration).build();
116+
}
117+
105118
public static final class DefaultBuilder implements QueryAsyncWaiter.Builder {
106119
private QueryAsyncClient client;
107120

‎codegen/src/test/resources/software/amazon/awssdk/codegen/poet/waiters/query-sync-waiter-class.java

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,13 @@
44
import java.util.ArrayList;
55
import java.util.List;
66
import java.util.Objects;
7+
import java.util.function.Consumer;
78
import software.amazon.awssdk.annotations.Generated;
89
import software.amazon.awssdk.annotations.SdkInternalApi;
910
import software.amazon.awssdk.annotations.ThreadSafe;
11+
import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
1012
import software.amazon.awssdk.awscore.exception.AwsServiceException;
13+
import software.amazon.awssdk.core.ApiName;
1114
import software.amazon.awssdk.core.internal.waiters.WaiterAttribute;
1215
import software.amazon.awssdk.core.retry.backoff.FixedDelayBackoffStrategy;
1316
import software.amazon.awssdk.core.waiters.PollingStrategy;
@@ -18,6 +21,7 @@
1821
import software.amazon.awssdk.services.query.QueryClient;
1922
import software.amazon.awssdk.services.query.model.APostOperationRequest;
2023
import software.amazon.awssdk.services.query.model.APostOperationResponse;
24+
import software.amazon.awssdk.services.query.model.QueryRequest;
2125
import software.amazon.awssdk.services.query.waiters.internal.WaitersRuntime;
2226
import software.amazon.awssdk.utils.AttributeMap;
2327
import software.amazon.awssdk.utils.SdkAutoCloseable;
@@ -59,7 +63,7 @@ private static String errorCode(Throwable error) {
5963

6064
@Override
6165
public WaiterResponse<APostOperationResponse> waitUntilPostOperationSuccess(APostOperationRequest aPostOperationRequest) {
62-
return postOperationSuccessWaiter.run(() -> client.aPostOperation(aPostOperationRequest));
66+
return postOperationSuccessWaiter.run(() -> client.aPostOperation(applyWaitersUserAgent(aPostOperationRequest)));
6367
}
6468

6569
private static List<WaiterAcceptor<? super APostOperationResponse>> postOperationSuccessWaiterAcceptors() {
@@ -84,6 +88,15 @@ public static QueryWaiter.Builder builder() {
8488
return new DefaultBuilder();
8589
}
8690

91+
private <T extends QueryRequest> T applyWaitersUserAgent(T request) {
92+
Consumer<AwsRequestOverrideConfiguration.Builder> userAgentApplier = b -> b.addApiName(ApiName.builder()
93+
.version("waiter").name("hll").build());
94+
AwsRequestOverrideConfiguration overrideConfiguration = request.overrideConfiguration()
95+
.map(c -> c.toBuilder().applyMutation(userAgentApplier).build())
96+
.orElse((AwsRequestOverrideConfiguration.builder().applyMutation(userAgentApplier).build()));
97+
return (T) request.toBuilder().overrideConfiguration(overrideConfiguration).build();
98+
}
99+
87100
public static final class DefaultBuilder implements QueryWaiter.Builder {
88101
private QueryClient client;
89102

‎services/autoscaling/src/it/java/software/amazon/awssdk/services/autoscaling/waiters/AutoScalingWaiterTest.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
package software.amazon.awssdk.services.autoscaling.waiters;
1717

1818
import static org.assertj.core.api.Assertions.assertThat;
19+
import static org.mockito.Matchers.any;
1920
import static org.mockito.Mockito.mock;
2021
import static org.mockito.Mockito.when;
2122
import static software.amazon.awssdk.services.autoscaling.model.LifecycleState.IN_SERVICE;
@@ -68,7 +69,7 @@ public void waitUntilGroupInServiceWorks() {
6869
i -> i.lifecycleState(IN_SERVICE)))
6970
.build();
7071

71-
when(client.describeAutoScalingGroups(request)).thenReturn(response1, response2);
72+
when(client.describeAutoScalingGroups(any(DescribeAutoScalingGroupsRequest.class))).thenReturn(response1, response2);
7273

7374
AutoScalingWaiter waiter = AutoScalingWaiter.builder()
7475
.pollingStrategy(PollingStrategy.builder()

‎services/dynamodb/src/test/java/software/amazon/awssdk/services/dynamodb/WaitersUserAgentTest.java

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License").
5+
* You may not use this file except in compliance with the License.
6+
* A copy of the License is located at
7+
*
8+
* http://aws.amazon.com/apache2.0
9+
*
10+
* or in the "license" file accompanying this file. This file is distributed
11+
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12+
* express or implied. See the License for the specific language governing
13+
* permissions and limitations under the License.
14+
*/
15+
16+
package software.amazon.awssdk.services.dynamodb;
17+
18+
import static com.github.tomakehurst.wiremock.client.WireMock.aResponse;
19+
import static com.github.tomakehurst.wiremock.client.WireMock.any;
20+
import static com.github.tomakehurst.wiremock.client.WireMock.stubFor;
21+
import static com.github.tomakehurst.wiremock.client.WireMock.urlEqualTo;
22+
import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.options;
23+
import static org.assertj.core.api.Assertions.assertThatThrownBy;
24+
import static org.junit.Assert.assertTrue;
25+
import static org.mockito.Mockito.times;
26+
27+
import com.github.tomakehurst.wiremock.common.ConsoleNotifier;
28+
import com.github.tomakehurst.wiremock.junit.WireMockRule;
29+
import java.net.URI;
30+
import java.util.concurrent.CompletableFuture;
31+
import org.junit.Before;
32+
import org.junit.Rule;
33+
import org.junit.Test;
34+
import org.mockito.ArgumentCaptor;
35+
import org.mockito.Matchers;
36+
import org.mockito.Mockito;
37+
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
38+
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
39+
import software.amazon.awssdk.core.interceptor.Context;
40+
import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
41+
import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
42+
import software.amazon.awssdk.core.waiters.WaiterResponse;
43+
import software.amazon.awssdk.regions.Region;
44+
import software.amazon.awssdk.services.dynamodb.model.DescribeTableRequest;
45+
import software.amazon.awssdk.services.dynamodb.model.DescribeTableResponse;
46+
import software.amazon.awssdk.services.dynamodb.waiters.DynamoDbAsyncWaiter;
47+
import software.amazon.awssdk.services.dynamodb.waiters.DynamoDbWaiter;
48+
49+
public class WaitersUserAgentTest {
50+
51+
@Rule
52+
public WireMockRule mockServer = new WireMockRule(options().notifier(new ConsoleNotifier(true)).dynamicHttpsPort());
53+
54+
private DynamoDbClient dynamoDbClient;
55+
private DynamoDbAsyncClient dynamoDbAsyncClient;
56+
private ExecutionInterceptor interceptor;
57+
58+
@Before
59+
public void setup() {
60+
interceptor = Mockito.spy(AbstractExecutionInterceptor.class);
61+
62+
dynamoDbClient = DynamoDbClient.builder()
63+
.credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("test",
64+
"test")))
65+
.region(Region.US_WEST_2)
66+
.endpointOverride(URI.create("http://localhost:" + mockServer.port()))
67+
.overrideConfiguration(c -> c.addExecutionInterceptor(interceptor))
68+
.build();
69+
70+
dynamoDbAsyncClient = DynamoDbAsyncClient.builder()
71+
.credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials
72+
.create("test",
73+
"test")))
74+
.region(Region.US_WEST_2)
75+
.endpointOverride(URI.create("http://localhost:" + mockServer.port()))
76+
.overrideConfiguration(c -> c.addExecutionInterceptor(interceptor))
77+
.build();
78+
}
79+
80+
@Test
81+
public void syncWaiters_shouldHaveWaitersUserAgent() {
82+
stubFor(any(urlEqualTo("/")).willReturn(aResponse().withStatus(500)));
83+
84+
DynamoDbWaiter waiter = dynamoDbClient.waiter();
85+
assertThatThrownBy(() -> waiter.waitUntilTableExists(DescribeTableRequest.builder().tableName("table").build())).isNotNull();
86+
87+
ArgumentCaptor<Context.BeforeTransmission> context = ArgumentCaptor.forClass(Context.BeforeTransmission.class);
88+
Mockito.verify(interceptor).beforeTransmission(context.capture(), Matchers.any());
89+
90+
assertTrue(context.getValue().httpRequest().headers().get("User-Agent").toString().contains("waiter"));
91+
}
92+
93+
@Test
94+
public void asyncWaiters_shouldHaveWaitersUserAgent() {
95+
DynamoDbAsyncWaiter waiter = dynamoDbAsyncClient.waiter();
96+
CompletableFuture<WaiterResponse<DescribeTableResponse>> responseFuture = waiter.waitUntilTableExists(DescribeTableRequest.builder().tableName("table").build());
97+
98+
ArgumentCaptor<Context.BeforeTransmission> context = ArgumentCaptor.forClass(Context.BeforeTransmission.class);
99+
Mockito.verify(interceptor).beforeTransmission(context.capture(), Matchers.any());
100+
101+
assertTrue(context.getValue().httpRequest().headers().get("User-Agent").toString().contains("waiter"));
102+
103+
responseFuture.cancel(true);
104+
}
105+
106+
public static abstract class AbstractExecutionInterceptor implements ExecutionInterceptor {
107+
@Override
108+
public void beforeTransmission(Context.BeforeTransmission context, ExecutionAttributes executionAttributes) {
109+
throw new RuntimeException("Interrupting the request.");
110+
}
111+
}
112+
}

‎services/ecs/src/it/java/software/amazon/awssdk/services/ecs/waiters/EcsWaiterTest.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
package software.amazon.awssdk.services.ecs.waiters;
1717

1818
import static org.assertj.core.api.Assertions.assertThat;
19+
import static org.mockito.Matchers.any;
1920
import static org.mockito.Mockito.mock;
2021
import static org.mockito.Mockito.when;
2122

@@ -57,7 +58,7 @@ public void waitUntilServicesStableWorks() {
5758
.runningCount(2))
5859
.build();
5960

60-
when(client.describeServices(request)).thenReturn(response1, response2);
61+
when(client.describeServices(any(DescribeServicesRequest.class))).thenReturn(response1, response2);
6162

6263
EcsWaiter waiter = EcsWaiter.builder()
6364
.pollingStrategy(PollingStrategy.builder()

0 commit comments

Comments
 (0)