Skip to content

Commit 305ba3a

Browse files
S-Saranya1S-Saranya1
committed
Move endpoint resolution from interceptors to pipeline stage
1 parent bfc51a7 commit 305ba3a

File tree

69 files changed

+3384
-1639
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

69 files changed

+3384
-1639
lines changed

‎codegen/src/main/java/software/amazon/awssdk/codegen/emitters/tasks/AuthSchemeGeneratorTasks.java‎

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,6 @@
2020
import software.amazon.awssdk.codegen.emitters.GeneratorTask;
2121
import software.amazon.awssdk.codegen.emitters.GeneratorTaskParams;
2222
import software.amazon.awssdk.codegen.emitters.PoetGeneratorTask;
23-
import software.amazon.awssdk.codegen.poet.auth.scheme.AuthSchemeInterceptorSpec;
2423
import software.amazon.awssdk.codegen.poet.auth.scheme.AuthSchemeParamsSpec;
2524
import software.amazon.awssdk.codegen.poet.auth.scheme.AuthSchemeProviderSpec;
2625
import software.amazon.awssdk.codegen.poet.auth.scheme.AuthSchemeSpecUtils;
@@ -85,10 +84,6 @@ private GeneratorTask generateEndpointAwareAuthSchemeParams() {
8584

8685
}
8786

88-
private GeneratorTask generateAuthSchemeInterceptor() {
89-
return new PoetGeneratorTask(authSchemeInternalDir(), model.getFileHeader(), new AuthSchemeInterceptorSpec(model));
90-
}
91-
9287
private String authSchemeDir() {
9388
return generatorTaskParams.getPathProvider().getAuthSchemeDirectory();
9489
}

‎codegen/src/main/java/software/amazon/awssdk/codegen/emitters/tasks/EndpointProviderTasks.java‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -32,9 +32,8 @@
3232
import software.amazon.awssdk.codegen.poet.rules.EndpointProviderInterfaceSpec;
3333
import software.amazon.awssdk.codegen.poet.rules.EndpointProviderSpec;
3434
import software.amazon.awssdk.codegen.poet.rules.EndpointProviderTestSpec;
35-
import software.amazon.awssdk.codegen.poet.rules.EndpointResolverInterceptorSpec;
35+
import software.amazon.awssdk.codegen.poet.rules.EndpointResolverUtilsSpec;
3636
import software.amazon.awssdk.codegen.poet.rules.EndpointRulesClientTestSpec;
37-
import software.amazon.awssdk.codegen.poet.rules.RequestEndpointInterceptorSpec;
3837
import software.amazon.awssdk.codegen.poet.rules2.EndpointProviderSpec2;
3938

4039
public final class EndpointProviderTasks extends BaseGeneratorTasks {
@@ -103,8 +102,7 @@ private boolean shouldGenerateCompiledEndpointRules() {
103102

104103
private Collection<GeneratorTask> generateInterceptors() {
105104
return Arrays.asList(
106-
new PoetGeneratorTask(endpointRulesInternalDir(), model.getFileHeader(), new EndpointResolverInterceptorSpec(model)),
107-
new PoetGeneratorTask(endpointRulesInternalDir(), model.getFileHeader(), new RequestEndpointInterceptorSpec(model)));
105+
new PoetGeneratorTask(endpointRulesInternalDir(), model.getFileHeader(), new EndpointResolverUtilsSpec(model)));
108106
}
109107

110108
private GeneratorTask generateClientTests() {

‎codegen/src/main/java/software/amazon/awssdk/codegen/poet/auth/scheme/AuthSchemeInterceptorSpec.java‎

Lines changed: 0 additions & 524 deletions
This file was deleted.

‎codegen/src/main/java/software/amazon/awssdk/codegen/poet/builder/BaseClientBuilderClass.java‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -370,9 +370,6 @@ private MethodSpec finalizeServiceConfigurationMethod() {
370370

371371
List<ClassName> builtInInterceptors = new ArrayList<>();
372372

373-
builtInInterceptors.add(endpointRulesSpecUtils.resolverInterceptorName());
374-
builtInInterceptors.add(endpointRulesSpecUtils.requestModifierInterceptorName());
375-
376373
for (String interceptor : model.getCustomizationConfig().getInterceptors()) {
377374
builtInInterceptors.add(ClassName.bestGuess(interceptor));
378375
}

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -172,7 +172,8 @@ protected void addAdditionalMethods(TypeSpec.Builder type) {
172172
.addMethods(protocolSpec.additionalMethods())
173173
.addMethod(protocolSpec.initProtocolFactory(model))
174174
.addMethod(resolveMetricPublishersMethod())
175-
.addMethod(ClientClassUtils.resolveAuthSchemeOptionsMethod(authSchemeSpecUtils, endpointRulesSpecUtils));
175+
.addMethod(ClientClassUtils.resolveAuthSchemeOptionsMethod(authSchemeSpecUtils, endpointRulesSpecUtils))
176+
.addMethod(ClientClassUtils.resolveEndpointMethod(authSchemeSpecUtils, endpointRulesSpecUtils));
176177

177178
type.addMethod(ClientClassUtils.updateRetryStrategyClientConfigurationMethod());
178179
type.addMethod(updateSdkClientConfigurationMethod(configurationUtils.serviceClientConfigurationBuilderClassName(),

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

Lines changed: 100 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -25,11 +25,13 @@
2525
import com.squareup.javapoet.ParameterizedTypeName;
2626
import com.squareup.javapoet.TypeName;
2727
import com.squareup.javapoet.TypeVariableName;
28+
import com.squareup.javapoet.WildcardTypeName;
2829
import java.util.List;
2930
import java.util.Map;
3031
import java.util.Objects;
3132
import java.util.Optional;
3233
import java.util.Set;
34+
import java.util.concurrent.CompletionException;
3335
import java.util.function.Consumer;
3436
import java.util.stream.Collectors;
3537
import javax.lang.model.element.Modifier;
@@ -51,11 +53,16 @@
5153
import software.amazon.awssdk.core.SdkClient;
5254
import software.amazon.awssdk.core.SdkPlugin;
5355
import software.amazon.awssdk.core.SdkRequest;
56+
import software.amazon.awssdk.core.SelectedAuthScheme;
5457
import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
5558
import software.amazon.awssdk.core.client.config.SdkClientConfiguration;
5659
import software.amazon.awssdk.core.client.config.SdkClientOption;
60+
import software.amazon.awssdk.core.exception.SdkClientException;
61+
import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
62+
import software.amazon.awssdk.core.interceptor.SdkInternalExecutionAttribute;
5763
import software.amazon.awssdk.core.retry.RetryMode;
5864
import software.amazon.awssdk.core.signer.Signer;
65+
import software.amazon.awssdk.endpoints.Endpoint;
5966
import software.amazon.awssdk.http.auth.spi.scheme.AuthSchemeOption;
6067
import software.amazon.awssdk.retries.api.RetryStrategy;
6168
import software.amazon.awssdk.utils.AttributeMap;
@@ -406,9 +413,6 @@ private static void addSimpleAuthSchemeResolution(MethodSpec.Builder builder,
406413
ClassName.get(Set.class), awsClientOption);
407414
builder.beginControlFlow("if (!$T.isNullOrEmpty(sigv4aRegionSet))", CollectionUtils.class);
408415
builder.addStatement("paramsBuilder.regionSet($T.create(sigv4aRegionSet))", regionSet);
409-
builder.nextControlFlow("else");
410-
builder.addStatement("paramsBuilder.regionSet($T.create(clientConfiguration.option($T.AWS_REGION).id()))",
411-
regionSet, awsClientOption);
412416
builder.endControlFlow();
413417
}
414418

@@ -422,7 +426,7 @@ private static void addEndpointBasedAuthSchemeResolution(MethodSpec.Builder buil
422426
ClassName paramsInterface = authSchemeSpecUtils.parametersInterfaceName();
423427
ClassName awsClientOption = ClassName.get("software.amazon.awssdk.awscore.client.config", "AwsClientOption");
424428
ClassName endpointParamsClass = endpointRulesSpecUtils.parametersClassName();
425-
ClassName resolverInterceptor = endpointRulesSpecUtils.resolverInterceptorName();
429+
ClassName endpointResolverUtils = endpointRulesSpecUtils.endpointResolverUtilsName();
426430
ClassName executionAttributesClass = ClassName.get("software.amazon.awssdk.core.interceptor", "ExecutionAttributes");
427431
ClassName awsExecutionAttribute = ClassName.get("software.amazon.awssdk.awscore", "AwsExecutionAttribute");
428432
ClassName sdkExecutionAttribute = ClassName.get("software.amazon.awssdk.core.interceptor", "SdkExecutionAttribute");
@@ -447,7 +451,7 @@ private static void addEndpointBasedAuthSchemeResolution(MethodSpec.Builder buil
447451
sdkInternalExecutionAttribute, SdkClientOption.class);
448452

449453
builder.addStatement("$T endpointParams = $T.ruleParams(request, executionAttributes)",
450-
endpointParamsClass, resolverInterceptor);
454+
endpointParamsClass, endpointResolverUtils);
451455

452456
builder.addStatement("$T.Builder paramsBuilder = $T.builder()", paramsInterface, paramsInterface);
453457

@@ -467,6 +471,15 @@ private static void addEndpointBasedAuthSchemeResolution(MethodSpec.Builder buil
467471
builder.addStatement("paramsBuilder.region(clientConfiguration.option($T.AWS_REGION))", awsClientOption);
468472
}
469473

474+
if (authSchemeSpecUtils.hasSigV4aSupport()) {
475+
ClassName regionSet = ClassName.get("software.amazon.awssdk.http.auth.aws.signer", "RegionSet");
476+
builder.addStatement("$T<String> sigv4aRegionSet = clientConfiguration.option($T.AWS_SIGV4A_SIGNING_REGION_SET)",
477+
ClassName.get(Set.class), awsClientOption);
478+
builder.beginControlFlow("if (!$T.isNullOrEmpty(sigv4aRegionSet))", CollectionUtils.class);
479+
builder.addStatement("paramsBuilder.regionSet($T.create(sigv4aRegionSet))", regionSet);
480+
builder.endControlFlow();
481+
}
482+
470483
ClassName paramsBuilderClass = authSchemeSpecUtils.parametersEndpointAwareDefaultImplName().nestedClass("Builder");
471484
ClassName endpointProviderInterface = endpointRulesSpecUtils.providerInterfaceName();
472485

@@ -482,4 +495,86 @@ private static void addEndpointBasedAuthSchemeResolution(MethodSpec.Builder buil
482495
builder.addStatement("$T<$T> options = authSchemeProvider.resolveAuthScheme(paramsBuilder.build())",
483496
List.class, AuthSchemeOption.class);
484497
}
498+
499+
static MethodSpec resolveEndpointMethod(AuthSchemeSpecUtils authSchemeSpecUtils,
500+
EndpointRulesSpecUtils endpointRulesSpecUtils) {
501+
ClassName utilsClass = endpointRulesSpecUtils.endpointResolverUtilsName();
502+
ClassName endpointParamsClass = endpointRulesSpecUtils.parametersClassName();
503+
ClassName providerInterface = endpointRulesSpecUtils.providerInterfaceName();
504+
ClassName awsEndpointProviderUtils = endpointRulesSpecUtils.sharedAwsEndpointProviderUtilsName();
505+
ClassName awsEndpointAttribute = ClassName.get("software.amazon.awssdk.awscore.endpoints", "AwsEndpointAttribute");
506+
ClassName endpointAuthScheme = ClassName.get("software.amazon.awssdk.awscore.endpoints.authscheme", "EndpointAuthScheme");
507+
508+
MethodSpec.Builder b = MethodSpec.methodBuilder("resolveEndpoint")
509+
.addModifiers(PRIVATE)
510+
.returns(Endpoint.class)
511+
.addParameter(SdkRequest.class, "request")
512+
.addParameter(ExecutionAttributes.class, "executionAttributes")
513+
.addParameter(String.class, "operationName");
514+
515+
b.addStatement("$1T provider = ($1T) executionAttributes.getAttribute($2T.ENDPOINT_PROVIDER)",
516+
providerInterface, SdkInternalExecutionAttribute.class);
517+
518+
b.beginControlFlow("try");
519+
b.addStatement("$T endpointParams = $T.ruleParams(request, executionAttributes)",
520+
endpointParamsClass, utilsClass);
521+
b.addStatement("$T endpoint = provider.resolveEndpoint(endpointParams).join()", Endpoint.class);
522+
523+
b.beginControlFlow("if (!$T.disableHostPrefixInjection(executionAttributes))", awsEndpointProviderUtils);
524+
b.addStatement("$T hostPrefix = $T.hostPrefix(operationName, request)",
525+
ParameterizedTypeName.get(Optional.class, String.class), utilsClass);
526+
b.beginControlFlow("if (hostPrefix.isPresent())");
527+
b.addStatement("endpoint = $T.addHostPrefix(endpoint, hostPrefix.get())", awsEndpointProviderUtils);
528+
b.endControlFlow();
529+
b.endControlFlow();
530+
531+
b.addStatement("$T endpointAuthSchemes = endpoint.attribute($T.AUTH_SCHEMES)",
532+
ParameterizedTypeName.get(ClassName.get(List.class), endpointAuthScheme),
533+
awsEndpointAttribute);
534+
b.addStatement("$T selectedAuthScheme = executionAttributes.getAttribute($T.SELECTED_AUTH_SCHEME)",
535+
ParameterizedTypeName.get(ClassName.get(SelectedAuthScheme.class),
536+
WildcardTypeName.subtypeOf(Object.class)),
537+
SdkInternalExecutionAttribute.class);
538+
b.beginControlFlow("if (endpointAuthSchemes != null && selectedAuthScheme != null)");
539+
b.addStatement("selectedAuthScheme = $T.authSchemeWithEndpointSignerProperties(endpointAuthSchemes, selectedAuthScheme)",
540+
utilsClass);
541+
542+
if (authSchemeSpecUtils.usesSigV4a() || authSchemeSpecUtils.generateEndpointBasedParams()) {
543+
ClassName awsV4aAuthScheme = ClassName.get("software.amazon.awssdk.http.auth.aws.scheme", "AwsV4aAuthScheme");
544+
ClassName awsV4aHttpSigner = ClassName.get("software.amazon.awssdk.http.auth.aws.signer", "AwsV4aHttpSigner");
545+
ClassName regionSet = ClassName.get("software.amazon.awssdk.http.auth.aws.signer", "RegionSet");
546+
ClassName authSchemeOption = ClassName.get("software.amazon.awssdk.http.auth.spi.scheme", "AuthSchemeOption");
547+
548+
b.addComment("Precedence of SigV4a RegionSet is set according to multi-auth SigV4a specifications");
549+
b.beginControlFlow("if (selectedAuthScheme.authSchemeOption().schemeId().equals($T.SCHEME_ID) "
550+
+ "&& selectedAuthScheme.authSchemeOption().signerProperty($T.REGION_SET) == null)",
551+
awsV4aAuthScheme, awsV4aHttpSigner);
552+
b.addStatement("$T optionBuilder = selectedAuthScheme.authSchemeOption().toBuilder()",
553+
authSchemeOption.nestedClass("Builder"));
554+
b.addStatement("$1T rs = $1T.create(endpointParams.region().id())", regionSet);
555+
b.addStatement("optionBuilder.putSignerProperty($T.REGION_SET, rs)", awsV4aHttpSigner);
556+
b.addStatement("selectedAuthScheme = new $T(selectedAuthScheme.identity(), selectedAuthScheme.signer(), "
557+
+ "optionBuilder.build())", SelectedAuthScheme.class);
558+
b.endControlFlow();
559+
}
560+
561+
b.addStatement("executionAttributes.putAttribute($T.SELECTED_AUTH_SCHEME, selectedAuthScheme)",
562+
SdkInternalExecutionAttribute.class);
563+
b.endControlFlow();
564+
565+
b.addStatement("$T.setMetricValues(endpoint, executionAttributes)", utilsClass);
566+
567+
b.addStatement("return endpoint");
568+
569+
b.nextControlFlow("catch ($T e)", CompletionException.class);
570+
b.addStatement("$T cause = e.getCause()", Throwable.class);
571+
b.beginControlFlow("if (cause instanceof $T)", SdkClientException.class);
572+
b.addStatement("throw ($T) cause", SdkClientException.class);
573+
b.endControlFlow();
574+
b.addStatement("throw $T.create($S + cause.getMessage(), cause)",
575+
SdkClientException.class, "Endpoint resolution failed: ");
576+
b.endControlFlow();
577+
578+
return b.build();
579+
}
485580
}

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,8 @@ protected void addAdditionalMethods(TypeSpec.Builder type) {
140140
.addMethod(nameMethod())
141141
.addMethods(protocolSpec.additionalMethods())
142142
.addMethod(resolveMetricPublishersMethod())
143-
.addMethod(ClientClassUtils.resolveAuthSchemeOptionsMethod(authSchemeSpecUtils, endpointRulesSpecUtils));
143+
.addMethod(ClientClassUtils.resolveAuthSchemeOptionsMethod(authSchemeSpecUtils, endpointRulesSpecUtils))
144+
.addMethod(ClientClassUtils.resolveEndpointMethod(authSchemeSpecUtils, endpointRulesSpecUtils));
144145

145146
protocolSpec.createErrorResponseHandler().ifPresent(type::addMethod);
146147
type.addMethod(ClientClassUtils.updateRetryStrategyClientConfigurationMethod());

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -224,6 +224,8 @@ public CodeBlock executionHandler(OperationModel opModel) {
224224
.add(".withMetricCollector(apiCallMetricCollector)\n")
225225
.add(".withAuthSchemeOptionsResolver(r -> resolveAuthSchemeOptions(r, $S, clientConfiguration))\n",
226226
opModel.getOperationName())
227+
.add(".withEndpointResolver((r, a) -> resolveEndpoint(r, a, $S))\n",
228+
opModel.getOperationName())
227229
.add(HttpChecksumRequiredTrait.putHttpChecksumAttribute(opModel))
228230
.add(HttpChecksumTrait.create(opModel));
229231

@@ -299,6 +301,8 @@ public CodeBlock asyncExecutionHandler(IntermediateModel intermediateModel, Oper
299301
.add(".withMetricCollector(apiCallMetricCollector)\n")
300302
.add(".withAuthSchemeOptionsResolver(r -> resolveAuthSchemeOptions(r, $S, clientConfiguration))\n",
301303
opModel.getOperationName())
304+
.add(".withEndpointResolver((r, a) -> resolveEndpoint(r, a, $S))\n",
305+
opModel.getOperationName())
302306
.add(hostPrefixExpression(opModel))
303307
.add(discoveredEndpoint(opModel))
304308
.add(credentialType(opModel, model))

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,8 @@ public CodeBlock executionHandler(OperationModel opModel) {
118118
.add(".withMetricCollector(apiCallMetricCollector)")
119119
.add(".withAuthSchemeOptionsResolver(r -> resolveAuthSchemeOptions(r, $S, clientConfiguration))\n",
120120
opModel.getOperationName())
121+
.add(".withEndpointResolver((r, a) -> resolveEndpoint(r, a, $S))\n",
122+
opModel.getOperationName())
121123
.add(HttpChecksumRequiredTrait.putHttpChecksumAttribute(opModel))
122124
.add(HttpChecksumTrait.create(opModel));
123125

@@ -159,6 +161,8 @@ public CodeBlock asyncExecutionHandler(IntermediateModel intermediateModel, Oper
159161
.add(".withMetricCollector(apiCallMetricCollector)\n")
160162
.add(".withAuthSchemeOptionsResolver(r -> resolveAuthSchemeOptions(r, $S, clientConfiguration))\n",
161163
opModel.getOperationName())
164+
.add(".withEndpointResolver((r, a) -> resolveEndpoint(r, a, $S))\n",
165+
opModel.getOperationName())
162166
.add(HttpChecksumRequiredTrait.putHttpChecksumAttribute(opModel))
163167
.add(HttpChecksumTrait.create(opModel));
164168

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -136,6 +136,8 @@ public CodeBlock executionHandler(OperationModel opModel) {
136136
.add(".withRequestConfiguration(clientConfiguration)")
137137
.add(".withInput($L)", opModel.getInput().getVariableName())
138138
.add(".withAuthSchemeOptionsResolver(r -> resolveAuthSchemeOptions(r, $S, clientConfiguration))\n",
139+
opModel.getOperationName())
140+
.add(".withEndpointResolver((r, a) -> resolveEndpoint(r, a, $S))\n",
139141
opModel.getOperationName())
140142
.add(HttpChecksumRequiredTrait.putHttpChecksumAttribute(opModel))
141143
.add(HttpChecksumTrait.create(opModel));
@@ -217,6 +219,8 @@ public CodeBlock asyncExecutionHandler(IntermediateModel intermediateModel, Oper
217219
.add(".withMetricCollector(apiCallMetricCollector)\n")
218220
.add(".withAuthSchemeOptionsResolver(r -> resolveAuthSchemeOptions(r, $S, clientConfiguration))\n",
219221
opModel.getOperationName())
222+
.add(".withEndpointResolver((r, a) -> resolveEndpoint(r, a, $S))\n",
223+
opModel.getOperationName())
220224
.add(asyncRequestBody(opModel))
221225
.add(HttpChecksumRequiredTrait.putHttpChecksumAttribute(opModel))
222226
.add(HttpChecksumTrait.create(opModel));

0 commit comments

Comments
 (0)