Skip to content

Commit 15242ce

Browse files
zoewanggzoewangg
committed
Reuse updateSdkClientConfigurationMethod in codegen for async and sync code paths so that the generated code is the same
1 parent 4166fa1 commit 15242ce

File tree

10 files changed

+602
-89
lines changed

10 files changed

+602
-89
lines changed

‎codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/AsyncClientClass.java‎

Lines changed: 3 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
import static software.amazon.awssdk.codegen.poet.client.ClientClassUtils.addS3ArnableFieldCode;
2828
import static software.amazon.awssdk.codegen.poet.client.ClientClassUtils.applySignerOverrideMethod;
2929
import static software.amazon.awssdk.codegen.poet.client.ClientClassUtils.transformServiceId;
30+
import static software.amazon.awssdk.codegen.poet.client.ClientClassUtils.updateSdkClientConfigurationMethod;
3031
import static software.amazon.awssdk.codegen.poet.client.SyncClientClass.addRequestModifierCode;
3132
import static software.amazon.awssdk.codegen.poet.client.SyncClientClass.getProtocolSpecs;
3233

@@ -77,8 +78,6 @@
7778
import software.amazon.awssdk.codegen.poet.model.EventStreamSpecHelper;
7879
import software.amazon.awssdk.codegen.poet.model.ServiceClientConfigurationUtils;
7980
import software.amazon.awssdk.core.RequestOverrideConfiguration;
80-
import software.amazon.awssdk.core.SdkPlugin;
81-
import software.amazon.awssdk.core.SdkRequest;
8281
import software.amazon.awssdk.core.async.AsyncResponseTransformer;
8382
import software.amazon.awssdk.core.async.AsyncResponseTransformerUtils;
8483
import software.amazon.awssdk.core.async.SdkPublisher;
@@ -181,7 +180,8 @@ protected void addAdditionalMethods(TypeSpec.Builder type) {
181180
}
182181
}
183182
type.addMethod(ClientClassUtils.updateRetryStrategyClientConfigurationMethod());
184-
type.addMethod(updateSdkClientConfigurationMethod(configurationUtils.serviceClientConfigurationBuilderClassName()));
183+
type.addMethod(updateSdkClientConfigurationMethod(configurationUtils.serviceClientConfigurationBuilderClassName(),
184+
model));
185185
protocolSpec.createErrorResponseHandler().ifPresent(type::addMethod);
186186
protocolSpec.createEventstreamErrorResponseHandler().ifPresent(type::addMethod);
187187
}
@@ -311,32 +311,6 @@ protected MethodSpec serviceClientConfigMethod() {
311311
.build();
312312
}
313313

314-
protected static MethodSpec updateSdkClientConfigurationMethod(
315-
TypeName serviceClientConfigurationBuilderClassName) {
316-
MethodSpec.Builder builder = MethodSpec.methodBuilder("updateSdkClientConfiguration")
317-
.addModifiers(PRIVATE)
318-
.addParameter(SdkRequest.class, "request")
319-
.addParameter(SdkClientConfiguration.class, "clientConfiguration")
320-
.returns(SdkClientConfiguration.class);
321-
322-
builder.addStatement("$T plugins = request.overrideConfiguration()\n"
323-
+ ".map(c -> c.plugins()).orElse(Collections.emptyList())",
324-
ParameterizedTypeName.get(List.class, SdkPlugin.class))
325-
.addStatement("$T configuration = clientConfiguration.toBuilder()", SdkClientConfiguration.Builder.class);
326-
327-
builder.beginControlFlow("if (plugins.isEmpty())")
328-
.addStatement("return configuration.build()")
329-
.endControlFlow()
330-
.addStatement("$1T serviceConfigBuilder = new $1T(configuration)", serviceClientConfigurationBuilderClassName)
331-
.beginControlFlow("for ($T plugin : plugins)", SdkPlugin.class)
332-
.addStatement("plugin.configureClient(serviceConfigBuilder)")
333-
.endControlFlow();
334-
builder.addStatement("updateRetryStrategyClientConfiguration(configuration)");
335-
builder.addStatement("return configuration.build()");
336-
337-
return builder.build();
338-
}
339-
340314
@Override
341315
protected void addCloseMethod(TypeSpec.Builder type) {
342316
MethodSpec method = MethodSpec.methodBuilder("close")

‎codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/ClientClassUtils.java‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515

1616
package software.amazon.awssdk.codegen.poet.client;
1717

18+
import static javax.lang.model.element.Modifier.PRIVATE;
1819
import static software.amazon.awssdk.codegen.poet.PoetUtils.classNameFromFqcn;
1920

2021
import com.squareup.javapoet.ClassName;
@@ -24,7 +25,9 @@
2425
import com.squareup.javapoet.ParameterizedTypeName;
2526
import com.squareup.javapoet.TypeName;
2627
import com.squareup.javapoet.TypeVariableName;
28+
import java.util.List;
2729
import java.util.Map;
30+
import java.util.Objects;
2831
import java.util.Optional;
2932
import java.util.function.Consumer;
3033
import java.util.stream.Collectors;
@@ -38,15 +41,21 @@
3841
import software.amazon.awssdk.codegen.model.intermediate.MemberModel;
3942
import software.amazon.awssdk.codegen.model.intermediate.OperationModel;
4043
import software.amazon.awssdk.codegen.model.intermediate.ShapeModel;
44+
import software.amazon.awssdk.codegen.model.service.ClientContextParam;
4145
import software.amazon.awssdk.codegen.model.service.HostPrefixProcessor;
4246
import software.amazon.awssdk.codegen.poet.PoetExtension;
4347
import software.amazon.awssdk.codegen.poet.PoetUtils;
48+
import software.amazon.awssdk.codegen.poet.rules.EndpointRulesSpecUtils;
49+
import software.amazon.awssdk.core.SdkPlugin;
50+
import software.amazon.awssdk.core.SdkRequest;
4451
import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
4552
import software.amazon.awssdk.core.client.config.SdkClientConfiguration;
4653
import software.amazon.awssdk.core.client.config.SdkClientOption;
4754
import software.amazon.awssdk.core.retry.RetryMode;
4855
import software.amazon.awssdk.core.signer.Signer;
4956
import software.amazon.awssdk.retries.api.RetryStrategy;
57+
import software.amazon.awssdk.utils.AttributeMap;
58+
import software.amazon.awssdk.utils.CollectionUtils;
5059
import software.amazon.awssdk.utils.HostnameValidator;
5160
import software.amazon.awssdk.utils.StringUtils;
5261
import software.amazon.awssdk.utils.Validate;
@@ -171,6 +180,57 @@ static CodeBlock addEndpointTraitCode(OperationModel opModel) {
171180
return builder.build();
172181
}
173182

183+
static MethodSpec updateSdkClientConfigurationMethod(
184+
TypeName serviceClientConfigurationBuilderClassName, IntermediateModel model) {
185+
MethodSpec.Builder builder = MethodSpec.methodBuilder("updateSdkClientConfiguration")
186+
.addModifiers(PRIVATE)
187+
.addParameter(SdkRequest.class, "request")
188+
.addParameter(SdkClientConfiguration.class, "clientConfiguration")
189+
.returns(SdkClientConfiguration.class);
190+
191+
builder.addStatement("$T plugins = request.overrideConfiguration()\n"
192+
+ ".map(c -> c.plugins()).orElse(Collections.emptyList())",
193+
ParameterizedTypeName.get(List.class, SdkPlugin.class))
194+
.addStatement("$T configuration = clientConfiguration.toBuilder()", SdkClientConfiguration.Builder.class);
195+
196+
builder.beginControlFlow("if (plugins.isEmpty())")
197+
.addStatement("return configuration.build()")
198+
.endControlFlow()
199+
.addStatement("$1T serviceConfigBuilder = new $1T(configuration)", serviceClientConfigurationBuilderClassName)
200+
.beginControlFlow("for ($T plugin : plugins)", SdkPlugin.class)
201+
.addStatement("plugin.configureClient(serviceConfigBuilder)")
202+
.endControlFlow();
203+
EndpointRulesSpecUtils endpointRulesSpecUtils = new EndpointRulesSpecUtils(model);
204+
205+
if (model.getCustomizationConfig() == null ||
206+
CollectionUtils.isNullOrEmpty(model.getCustomizationConfig().getCustomClientContextParams())) {
207+
builder.addStatement("updateRetryStrategyClientConfiguration(configuration)");
208+
builder.addStatement("return configuration.build()");
209+
return builder.build();
210+
}
211+
212+
Map<String, ClientContextParam> customClientConfigParams = model.getCustomizationConfig().getCustomClientContextParams();
213+
214+
builder.addCode("$1T newContextParams = configuration.option($2T.CLIENT_CONTEXT_PARAMS);\n"
215+
+ "$1T originalContextParams = clientConfiguration.option($2T.CLIENT_CONTEXT_PARAMS);",
216+
AttributeMap.class, SdkClientOption.class);
217+
218+
builder.addCode("newContextParams = (newContextParams != null) ? newContextParams : $1T.empty();\n"
219+
+ "originalContextParams = originalContextParams != null ? originalContextParams : $1T.empty();",
220+
AttributeMap.class);
221+
222+
customClientConfigParams.forEach((n, m) -> {
223+
String keyName = model.getNamingStrategy().getEnumValueName(n);
224+
builder.addStatement("$1T.validState($2T.equals(originalContextParams.get($3T.$4N), newContextParams.get($3T.$4N)),"
225+
+ " $5S)",
226+
Validate.class, Objects.class, endpointRulesSpecUtils.clientContextParamsName(), keyName,
227+
keyName + " cannot be modified by request level plugins");
228+
});
229+
builder.addStatement("updateRetryStrategyClientConfiguration(configuration)");
230+
builder.addStatement("return configuration.build()");
231+
return builder.build();
232+
}
233+
174234
static Optional<CodeBlock> addS3ArnableFieldCode(OperationModel opModel, IntermediateModel model) {
175235
CodeBlock.Builder builder = CodeBlock.builder();
176236
Map<String, S3ArnableFieldConfig> s3ArnableFields = model.getCustomizationConfig().getS3ArnableFields();

‎codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/SyncClientClass.java‎

Lines changed: 3 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -24,21 +24,20 @@
2424
import static software.amazon.awssdk.codegen.poet.client.ClientClassUtils.addS3ArnableFieldCode;
2525
import static software.amazon.awssdk.codegen.poet.client.ClientClassUtils.applySignerOverrideMethod;
2626
import static software.amazon.awssdk.codegen.poet.client.ClientClassUtils.transformServiceId;
27+
import static software.amazon.awssdk.codegen.poet.client.ClientClassUtils.updateSdkClientConfigurationMethod;
2728

2829
import com.squareup.javapoet.ClassName;
2930
import com.squareup.javapoet.CodeBlock;
3031
import com.squareup.javapoet.FieldSpec;
3132
import com.squareup.javapoet.MethodSpec;
3233
import com.squareup.javapoet.ParameterizedTypeName;
33-
import com.squareup.javapoet.TypeName;
3434
import com.squareup.javapoet.TypeSpec;
3535
import com.squareup.javapoet.WildcardTypeName;
3636
import java.net.URI;
3737
import java.util.ArrayList;
3838
import java.util.Collections;
3939
import java.util.List;
4040
import java.util.Map;
41-
import java.util.Objects;
4241
import java.util.Optional;
4342
import java.util.concurrent.CompletableFuture;
4443
import java.util.stream.Collectors;
@@ -53,7 +52,6 @@
5352
import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
5453
import software.amazon.awssdk.codegen.model.intermediate.OperationModel;
5554
import software.amazon.awssdk.codegen.model.intermediate.Protocol;
56-
import software.amazon.awssdk.codegen.model.service.ClientContextParam;
5755
import software.amazon.awssdk.codegen.model.service.PreClientExecutionRequestCustomizer;
5856
import software.amazon.awssdk.codegen.poet.PoetExtension;
5957
import software.amazon.awssdk.codegen.poet.PoetUtils;
@@ -64,10 +62,7 @@
6462
import software.amazon.awssdk.codegen.poet.client.specs.QueryProtocolSpec;
6563
import software.amazon.awssdk.codegen.poet.client.specs.XmlProtocolSpec;
6664
import software.amazon.awssdk.codegen.poet.model.ServiceClientConfigurationUtils;
67-
import software.amazon.awssdk.codegen.poet.rules.EndpointRulesSpecUtils;
6865
import software.amazon.awssdk.core.RequestOverrideConfiguration;
69-
import software.amazon.awssdk.core.SdkPlugin;
70-
import software.amazon.awssdk.core.SdkRequest;
7166
import software.amazon.awssdk.core.client.config.SdkClientConfiguration;
7267
import software.amazon.awssdk.core.client.config.SdkClientOption;
7368
import software.amazon.awssdk.core.client.handler.SyncClientHandler;
@@ -78,11 +73,9 @@
7873
import software.amazon.awssdk.metrics.MetricCollector;
7974
import software.amazon.awssdk.metrics.MetricPublisher;
8075
import software.amazon.awssdk.metrics.NoOpMetricCollector;
81-
import software.amazon.awssdk.utils.AttributeMap;
8276
import software.amazon.awssdk.utils.CollectionUtils;
8377
import software.amazon.awssdk.utils.CompletableFutureUtils;
8478
import software.amazon.awssdk.utils.Logger;
85-
import software.amazon.awssdk.utils.Validate;
8679

8780
public class SyncClientClass extends SyncClientInterface {
8881

@@ -152,7 +145,8 @@ protected void addAdditionalMethods(TypeSpec.Builder type) {
152145

153146
protocolSpec.createErrorResponseHandler().ifPresent(type::addMethod);
154147
type.addMethod(ClientClassUtils.updateRetryStrategyClientConfigurationMethod());
155-
type.addMethod(updateSdkClientConfigurationMethod(configurationUtils.serviceClientConfigurationBuilderClassName()));
148+
type.addMethod(updateSdkClientConfigurationMethod(configurationUtils.serviceClientConfigurationBuilderClassName(),
149+
model));
156150
type.addMethod(protocolSpec.initProtocolFactory(model));
157151
}
158152

@@ -462,55 +456,4 @@ protected MethodSpec.Builder waiterOperationBody(MethodSpec.Builder builder) {
462456
.addStatement("return $T.builder().client(this).build()",
463457
poetExtensions.getSyncWaiterInterface());
464458
}
465-
466-
protected MethodSpec updateSdkClientConfigurationMethod(
467-
TypeName serviceClientConfigurationBuilderClassName) {
468-
MethodSpec.Builder builder = MethodSpec.methodBuilder("updateSdkClientConfiguration")
469-
.addModifiers(PRIVATE)
470-
.addParameter(SdkRequest.class, "request")
471-
.addParameter(SdkClientConfiguration.class, "clientConfiguration")
472-
.returns(SdkClientConfiguration.class);
473-
474-
builder.addStatement("$T plugins = request.overrideConfiguration()\n"
475-
+ ".map(c -> c.plugins()).orElse(Collections.emptyList())",
476-
ParameterizedTypeName.get(List.class, SdkPlugin.class))
477-
.addStatement("$T configuration = clientConfiguration.toBuilder()", SdkClientConfiguration.Builder.class);
478-
479-
builder.beginControlFlow("if (plugins.isEmpty())")
480-
.addStatement("return configuration.build()")
481-
.endControlFlow()
482-
.addStatement("$1T serviceConfigBuilder = new $1T(configuration)", serviceClientConfigurationBuilderClassName)
483-
.beginControlFlow("for ($T plugin : plugins)", SdkPlugin.class)
484-
.addStatement("plugin.configureClient(serviceConfigBuilder)")
485-
.endControlFlow();
486-
EndpointRulesSpecUtils endpointRulesSpecUtils = new EndpointRulesSpecUtils(this.model);
487-
488-
if (model.getCustomizationConfig() == null ||
489-
CollectionUtils.isNullOrEmpty(model.getCustomizationConfig().getCustomClientContextParams())) {
490-
builder.addStatement("updateRetryStrategyClientConfiguration(configuration)");
491-
builder.addStatement("return configuration.build()");
492-
return builder.build();
493-
}
494-
495-
Map<String, ClientContextParam> customClientConfigParams = model.getCustomizationConfig().getCustomClientContextParams();
496-
497-
builder.addCode("$1T newContextParams = configuration.option($2T.CLIENT_CONTEXT_PARAMS);\n"
498-
+ "$1T originalContextParams = clientConfiguration.option($2T.CLIENT_CONTEXT_PARAMS);",
499-
AttributeMap.class, SdkClientOption.class);
500-
501-
builder.addCode("newContextParams = (newContextParams != null) ? newContextParams : $1T.empty();\n"
502-
+ "originalContextParams = originalContextParams != null ? originalContextParams : $1T.empty();",
503-
AttributeMap.class);
504-
505-
customClientConfigParams.forEach((n, m) -> {
506-
String keyName = model.getNamingStrategy().getEnumValueName(n);
507-
builder.addStatement("$1T.validState($2T.equals(originalContextParams.get($3T.$4N), newContextParams.get($3T.$4N)),"
508-
+ " $5S)",
509-
Validate.class, Objects.class, endpointRulesSpecUtils.clientContextParamsName(), keyName,
510-
keyName + " cannot be modified by request level plugins");
511-
});
512-
builder.addStatement("updateRetryStrategyClientConfiguration(configuration)");
513-
builder.addStatement("return configuration.build()");
514-
return builder.build();
515-
}
516459
}

‎codegen/src/test/java/software/amazon/awssdk/codegen/poet/ClientTestModels.java‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -427,6 +427,21 @@ public static IntermediateModel opsWithSigv4a() {
427427
return new IntermediateModelBuilder(models).build();
428428
}
429429

430+
public static IntermediateModel pluginModels() {
431+
File serviceModel =
432+
new File(ClientTestModels.class.getResource("client/c2j/service-with-custom-context-params/service-2.json").getFile());
433+
File customizationModel =
434+
new File(ClientTestModels.class.getResource("client/c2j/service-with-custom-context-params/customization.config")
435+
.getFile());
436+
C2jModels models = C2jModels
437+
.builder()
438+
.serviceModel(getServiceModel(serviceModel))
439+
.customizationConfig(getCustomizationConfig(customizationModel))
440+
.build();
441+
442+
return new IntermediateModelBuilder(models).build();
443+
}
444+
430445
public static IntermediateModel xmlServiceModels() {
431446
File serviceModel = new File(ClientTestModels.class.getResource("client/c2j/xml/service-2.json").getFile());
432447
File customizationModel = new File(ClientTestModels.class.getResource("client/c2j/xml/customization.config").getFile());

‎codegen/src/test/java/software/amazon/awssdk/codegen/poet/client/AsyncClientClassTest.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import static software.amazon.awssdk.codegen.poet.ClientTestModels.customPackageModels;
2525
import static software.amazon.awssdk.codegen.poet.ClientTestModels.endpointDiscoveryModels;
2626
import static software.amazon.awssdk.codegen.poet.ClientTestModels.opsWithSigv4a;
27+
import static software.amazon.awssdk.codegen.poet.ClientTestModels.pluginModels;
2728
import static software.amazon.awssdk.codegen.poet.ClientTestModels.queryServiceModels;
2829
import static software.amazon.awssdk.codegen.poet.ClientTestModels.restJsonServiceModels;
2930
import static software.amazon.awssdk.codegen.poet.ClientTestModels.rpcv2ServiceModels;
@@ -123,6 +124,12 @@ public void asyncClientWithStreamingUnsignedPayload() {
123124
assertThat(asyncClientClass, generatesTo("test-unsigned-payload-trait-async-client-class.java"));
124125
}
125126

127+
@Test
128+
public void asyncClientWithPlugin() {
129+
AsyncClientClass asyncClientClass = createAsyncClientClass(pluginModels());
130+
assertThat(asyncClientClass, generatesTo("test-custom-context-params-async-client-class.java"));
131+
}
132+
126133
private AsyncClientClass createAsyncClientClass(IntermediateModel model) {
127134
return new AsyncClientClass(GeneratorTaskParams.create(model, "sources/", "tests/", "resources/"));
128135
}

‎codegen/src/test/java/software/amazon/awssdk/codegen/poet/client/SyncClientClassTest.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import static software.amazon.awssdk.codegen.poet.ClientTestModels.customPackageModels;
2323
import static software.amazon.awssdk.codegen.poet.ClientTestModels.endpointDiscoveryModels;
2424
import static software.amazon.awssdk.codegen.poet.ClientTestModels.opsWithSigv4a;
25+
import static software.amazon.awssdk.codegen.poet.ClientTestModels.pluginModels;
2526
import static software.amazon.awssdk.codegen.poet.ClientTestModels.queryServiceModels;
2627
import static software.amazon.awssdk.codegen.poet.ClientTestModels.restJsonServiceModels;
2728
import static software.amazon.awssdk.codegen.poet.ClientTestModels.rpcv2ServiceModels;
@@ -97,6 +98,12 @@ public void syncClientClassCbor() {
9798
assertThat(syncClientCustomServiceMetaData, generatesTo("test-cbor-client-class.java"));
9899
}
99100

101+
@Test
102+
public void asyncClientWithCustomContextParams() {
103+
ClassSpec syncClientClass = createSyncClientClass(pluginModels());
104+
assertThat(syncClientClass, generatesTo("test-custom-context-params-sync-client-class.java"));
105+
}
106+
100107
private SyncClientClass createSyncClientClass(IntermediateModel model) {
101108
return new SyncClientClass(GeneratorTaskParams.create(model, "sources/", "tests/", "resources/"));
102109
}

‎codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/c2j/service-with-custom-context-params/customization.config‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
{
2+
"customClientContextParams":{
3+
"CrossRegionAccessEnabled":{
4+
"documentation":"Enables cross-region access for this client",
5+
"type":"boolean"
6+
}
7+
}
8+
}

0 commit comments

Comments
 (0)