Skip to content

Commit

Permalink
Add AuthScheme to the SerivceClientConfiguration
Browse files Browse the repository at this point in the history
  • Loading branch information
sugmanue committed Oct 6, 2023
1 parent e064ef0 commit be4d511
Show file tree
Hide file tree
Showing 6 changed files with 233 additions and 4 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -165,6 +165,10 @@ private void addLocalFieldForBuilderIfNeeded(Field field, TypeSpec.Builder build
}

private MethodSpec setterForField(Field field) {
MethodSpec fieldBuilderSetter = field.builderSetterImpl();
if (fieldBuilderSetter != null) {
return fieldBuilderSetter.toBuilder().returns(builderInterface).build();
}
MethodSpec.Builder builder = baseSetterForField(field);
if (field.isLocalField()) {
builder.addAnnotation(Override.class);
Expand Down Expand Up @@ -199,6 +203,12 @@ private MethodSpec getterForBuilderField(Field field) {


private MethodSpec getterForField(Field field, String fieldName, boolean forDataGetter) {
MethodSpec fieldBuilderGetter = field.builderGetterImpl();
if (fieldBuilderGetter != null) {
return fieldBuilderGetter.toBuilder()
.returns(field.type())
.build();
}
MethodSpec.Builder builder = baseGetterForField(field);
if (!forDataGetter && field.isLocalField()) {
builder.addAnnotation(Override.class);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -122,6 +122,12 @@ private void addLocalFieldForDataIfNeeded(Field field, TypeSpec.Builder builder)
}

private MethodSpec.Builder baseSetterForField(Field field) {
MethodSpec fieldBuilderSetter = field.builderSetter();
if (fieldBuilderSetter != null) {
return fieldBuilderSetter.toBuilder()
.returns(className().nestedClass("Builder"));
}

MethodSpec.Builder builder = MethodSpec.methodBuilder(field.name())
.addModifiers(PUBLIC)
.addParameter(field.type(), field.name())
Expand All @@ -138,6 +144,13 @@ private MethodSpec getterForDataField(Field field) {
}

private MethodSpec getterForField(Field field, String fieldName, boolean forDataGetter) {
MethodSpec fieldBuilderGetter = field.builderGetterImpl();
if (fieldBuilderGetter != null) {
return fieldBuilderGetter.toBuilder()
.returns(field.type())
.build();
}

MethodSpec.Builder builder = baseGetterForField(field);
if (!forDataGetter && field.isLocalField()) {
builder.addAnnotation(Override.class);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,13 +17,18 @@

import com.squareup.javapoet.ClassName;
import com.squareup.javapoet.CodeBlock;
import com.squareup.javapoet.MethodSpec;
import com.squareup.javapoet.ParameterizedTypeName;
import com.squareup.javapoet.TypeName;
import com.squareup.javapoet.WildcardTypeName;
import java.net.URI;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import javax.lang.model.element.Modifier;
import software.amazon.awssdk.awscore.AwsServiceClientConfiguration;
import software.amazon.awssdk.awscore.client.config.AwsClientOption;
import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
Expand All @@ -34,6 +39,7 @@
import software.amazon.awssdk.core.client.config.SdkClientConfiguration;
import software.amazon.awssdk.core.client.config.SdkClientOption;
import software.amazon.awssdk.endpoints.EndpointProvider;
import software.amazon.awssdk.http.auth.spi.scheme.AuthScheme;
import software.amazon.awssdk.http.auth.spi.scheme.AuthSchemeProvider;
import software.amazon.awssdk.identity.spi.AwsCredentialsIdentity;
import software.amazon.awssdk.identity.spi.IdentityProvider;
Expand Down Expand Up @@ -132,7 +138,8 @@ private static List<Field> baseServiceClientConfigurationFields() {
.optionClass(AwsClientOption.class)
.optionValue(AwsClientOption.AWS_REGION)
.build(),
credentialsProviderField()
credentialsProviderField(),
authSchemesField()
);
}

Expand Down Expand Up @@ -215,6 +222,78 @@ SdkClientOption.class, fieldName(SdkClientOption.IDENTITY_PROVIDERS, SdkClientOp
return builder.build();
}


private static Field authSchemesField() {
TypeName authSchemeGenericType = ParameterizedTypeName.get(ClassName.get(AuthScheme.class),
WildcardTypeName.subtypeOf(Object.class));
TypeName authSchemesType = ParameterizedTypeName.get(ClassName.get(Map.class), ClassName.get(String.class),
authSchemeGenericType);
Field.Builder builder = Field.builder("authSchemes",
authSchemesType)
.doc("auth schemes")
.definingClass(AwsServiceClientConfiguration.class);

builder.constructFromConfiguration(
CodeBlock.builder()
.addStatement("$T authSchemes = internalBuilder.option($T.$L)",
authSchemesType, SdkClientOption.class,
fieldName(SdkClientOption.AUTH_SCHEMES, SdkClientOption.class))
.beginControlFlow("if (authSchemes != null)")
.addStatement("authSchemes = new $T<>(authSchemes)", HashMap.class)
.endControlFlow()
.addStatement("this.authSchemes = authSchemes")
.build()
);

builder.copyToConfiguration(
CodeBlock.builder()
.beginControlFlow("if (authSchemes != null &&"
+ " !authSchemes.equals(internalBuilder.option($T.$L)))",
SdkClientOption.class, fieldName(SdkClientOption.AUTH_SCHEMES,
SdkClientOption.class))
.addStatement("internalBuilder.option($T.$L, authSchemes())",
SdkClientOption.class, fieldName(SdkClientOption.AUTH_SCHEMES,
SdkClientOption.class))
.endControlFlow()
.build()
);

builder.builderSetterImpl(
MethodSpec.methodBuilder("putAuthScheme")
.addModifiers(Modifier.PUBLIC)
.addAnnotation(Override.class)
.addParameter(authSchemeGenericType, "authScheme")
.beginControlFlow("if (authSchemes == null)")
.addStatement("authSchemes = new $T<>()", HashMap.class)
.endControlFlow()
.addStatement("authSchemes.put(authScheme.schemeId(), authScheme)")
.addStatement("return this")
.build()
);

builder.builderSetter(
MethodSpec.methodBuilder("putAuthScheme")
.addModifiers(Modifier.PUBLIC)
.addAnnotation(Override.class)
.addParameter(authSchemeGenericType, "authScheme")
.build()
);

builder.builderGetterImpl(
MethodSpec.methodBuilder("authSchemes")
.addModifiers(Modifier.PUBLIC)
.addAnnotation(Override.class)
.returns(authSchemesType)
.beginControlFlow("if (authSchemes == null)")
.addStatement("return $T.emptyMap()", Collections.class)
.endControlFlow()
.addStatement("return $T.unmodifiableMap(new $T<>(authSchemes))", Collections.class, HashMap.class)
.build()
);

return builder.build();
}

static class Field {
private final String name;
private final TypeName type;
Expand All @@ -225,6 +304,9 @@ static class Field {
private final TypeName baseType;
private final CodeBlock constructFromConfiguration;
private final CodeBlock copyToConfiguration;
private final MethodSpec builderSetterImpl;
private final MethodSpec builderSetter;
private final MethodSpec builderGetterImpl;

Field(Field.Builder builder) {
this.name = Validate.paramNotNull(builder.name, "name");
Expand All @@ -236,6 +318,9 @@ static class Field {
this.baseType = builder.baseType;
this.constructFromConfiguration = builder.constructFromConfiguration;
this.copyToConfiguration = builder.copyToConfiguration;
this.builderSetterImpl = builder.builderSetterImpl;
this.builderSetter = builder.builderSetter;
this.builderGetterImpl = builder.builderGetterImpl;
}

public boolean isLocalField() {
Expand Down Expand Up @@ -278,6 +363,18 @@ public CodeBlock copyToConfiguration() {
return copyToConfiguration;
}

public MethodSpec builderSetterImpl() {
return builderSetterImpl;
}

public MethodSpec builderSetter() {
return builderSetter;
}

public MethodSpec builderGetterImpl() {
return builderGetterImpl;
}

public static Field.Builder builder() {
return new Field.Builder();
}
Expand Down Expand Up @@ -305,7 +402,9 @@ static class Builder {
private TypeName baseType;
private CodeBlock constructFromConfiguration;
private CodeBlock copyToConfiguration;

private MethodSpec builderSetterImpl;
private MethodSpec builderSetter;
private MethodSpec builderGetterImpl;

public Field.Builder name(String name) {
this.name = name;
Expand Down Expand Up @@ -357,6 +456,21 @@ public Field.Builder copyToConfiguration(CodeBlock copyToConfiguration) {
return this;
}

public Field.Builder builderSetterImpl(MethodSpec builderSetter) {
this.builderSetterImpl = builderSetter;
return this;
}

public Field.Builder builderSetter(MethodSpec builderSetter) {
this.builderSetter = builderSetter;
return this;
}

public Field.Builder builderGetterImpl(MethodSpec builderGetterImpl) {
this.builderGetterImpl = builderGetterImpl;
return this;
}

public Field build() {
if (value != null && optionClass != null) {
optionName = fieldName(value, optionClass);
Expand Down
Original file line number Diff line number Diff line change
@@ -1,13 +1,17 @@
package software.amazon.awssdk.services.jsonprotocoltests.internal;

import java.net.URI;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import software.amazon.awssdk.annotations.Generated;
import software.amazon.awssdk.annotations.SdkInternalApi;
import software.amazon.awssdk.awscore.client.config.AwsClientOption;
import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
import software.amazon.awssdk.core.client.config.SdkClientConfiguration;
import software.amazon.awssdk.core.client.config.SdkClientOption;
import software.amazon.awssdk.endpoints.EndpointProvider;
import software.amazon.awssdk.http.auth.spi.scheme.AuthScheme;
import software.amazon.awssdk.http.auth.spi.scheme.AuthSchemeProvider;
import software.amazon.awssdk.identity.spi.AwsCredentialsIdentity;
import software.amazon.awssdk.identity.spi.IdentityProvider;
Expand Down Expand Up @@ -41,6 +45,8 @@ public static class BuilderImpl implements BuilderInternal {

private IdentityProvider<? extends AwsCredentialsIdentity> credentialsProvider;

private Map<String, AuthScheme<?>> authSchemes;

private BuilderImpl() {
this.internalBuilder = SdkClientConfiguration.builder();
}
Expand All @@ -53,6 +59,11 @@ private BuilderImpl(SdkClientConfiguration.Builder internalBuilder) {
this.endpointOverride = internalBuilder.option(SdkClientOption.ENDPOINT);
}
this.credentialsProvider = internalBuilder.option(AwsClientOption.CREDENTIALS_IDENTITY_PROVIDER);
Map<String, AuthScheme<?>> authSchemes = internalBuilder.option(SdkClientOption.AUTH_SCHEMES);
if (authSchemes != null) {
authSchemes = new HashMap<>(authSchemes);
}
this.authSchemes = authSchemes;
}

/**
Expand Down Expand Up @@ -142,6 +153,23 @@ public IdentityProvider<? extends AwsCredentialsIdentity> credentialsProvider()
return credentialsProvider;
}

@Override
public JsonProtocolTestsServiceClientConfiguration.Builder putAuthScheme(AuthScheme<?> authScheme) {
if (authSchemes == null) {
authSchemes = new HashMap<>();
}
authSchemes.put(authScheme.schemeId(), authScheme);
return this;
}

@Override
public Map<String, AuthScheme<?>> authSchemes() {
if (authSchemes == null) {
return Collections.emptyMap();
}
return Collections.unmodifiableMap(new HashMap<>(authSchemes));
}

/**
* Sets the value for auth scheme provider
*/
Expand Down Expand Up @@ -190,6 +218,9 @@ public SdkClientConfiguration buildSdkClientConfiguration() {
}
internalBuilder.option(SdkClientOption.IDENTITY_PROVIDERS, identityProviders);
}
if (authSchemes != null && !authSchemes.equals(internalBuilder.option(SdkClientOption.AUTH_SCHEMES))) {
internalBuilder.option(SdkClientOption.AUTH_SCHEMES, authSchemes());
}
return internalBuilder.build();
}
}
Expand Down
Original file line number Diff line number Diff line change
@@ -1,11 +1,13 @@
package software.amazon.awssdk.services.jsonprotocoltests;

import java.net.URI;
import java.util.Map;
import software.amazon.awssdk.annotations.Generated;
import software.amazon.awssdk.annotations.SdkPublicApi;
import software.amazon.awssdk.awscore.AwsServiceClientConfiguration;
import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
import software.amazon.awssdk.endpoints.EndpointProvider;
import software.amazon.awssdk.http.auth.spi.scheme.AuthScheme;
import software.amazon.awssdk.identity.spi.AwsCredentialsIdentity;
import software.amazon.awssdk.identity.spi.IdentityProvider;
import software.amazon.awssdk.regions.Region;
Expand Down Expand Up @@ -100,6 +102,15 @@ public interface Builder extends AwsServiceClientConfiguration.Builder {
@Override
IdentityProvider<? extends AwsCredentialsIdentity> credentialsProvider();

@Override
Builder putAuthScheme(AuthScheme<?> authScheme);

/**
* Gets the value for auth schemes
*/
@Override
Map<String, AuthScheme<?>> authSchemes();

/**
* Sets the value for auth scheme provider
*/
Expand Down
Loading

0 comments on commit be4d511

Please sign in to comment.