Skip to content

Codegen the partitions data #4789

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Jan 11, 2024
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -19,31 +19,49 @@
import com.squareup.javapoet.CodeBlock;
import com.squareup.javapoet.FieldSpec;
import com.squareup.javapoet.MethodSpec;
import com.squareup.javapoet.ParameterizedTypeName;
import com.squareup.javapoet.TypeSpec;
import java.io.IOException;
import java.io.InputStream;
import java.io.UncheckedIOException;
import java.util.Map;
import javax.lang.model.element.Modifier;
import software.amazon.awssdk.annotations.SdkInternalApi;
import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
import software.amazon.awssdk.codegen.poet.ClassSpec;
import software.amazon.awssdk.codegen.poet.PoetUtils;
import software.amazon.awssdk.protocols.jsoncore.JsonNode;
import software.amazon.awssdk.utils.IoUtils;
import software.amazon.awssdk.utils.Lazy;
import software.amazon.awssdk.utils.Validate;

public class DefaultPartitionDataProviderSpec implements ClassSpec {
// partitions
private static final String VERSION = "version";
private static final String PARTITIONS = "partitions";
// partition
private static final String ID = "id";
private static final String REGION_REGEX = "regionRegex";
private static final String REGIONS = "regions";
private static final String OUTPUTS = "outputs";
// outputs
private static final String DNS_SUFFIX = "dnsSuffix";
private static final String DUAL_STACK_DNS_SUFFIX = "dualStackDnsSuffix";
private static final String SUPPORTS_FIPS = "supportsFIPS";
private static final String SUPPORTS_DUAL_STACK = "supportsDualStack";

private final IntermediateModel model;
private final EndpointRulesSpecUtils endpointRulesSpecUtils;

private final ClassName partitionsClass;
private final ClassName partitionClass;
private final ClassName regionOverrideClass;
private final ClassName outputsClass;

public DefaultPartitionDataProviderSpec(IntermediateModel model) {
this.model = model;
this.endpointRulesSpecUtils = new EndpointRulesSpecUtils(model);
this.partitionsClass = endpointRulesSpecUtils.rulesRuntimeClassName("Partitions");
this.partitionClass = endpointRulesSpecUtils.rulesRuntimeClassName("Partition");
this.regionOverrideClass = endpointRulesSpecUtils.rulesRuntimeClassName("RegionOverride");
this.outputsClass = endpointRulesSpecUtils.rulesRuntimeClassName("Outputs");
}

@Override
Expand All @@ -54,10 +72,8 @@ public TypeSpec poetSpec() {
.addSuperinterface(
endpointRulesSpecUtils.rulesRuntimeClassName("PartitionDataProvider"));

builder.addField(partitionDataField());
builder.addField(partitionsLazyField());
builder.addType(lazyPartitionsContainer());
builder.addMethod(loadPartitionsMethod());
builder.addMethod(doLoadPartitionsMethod());
return builder.build();
}

Expand All @@ -68,44 +84,24 @@ public ClassName className() {

private MethodSpec loadPartitionsMethod() {
MethodSpec.Builder builder = MethodSpec.methodBuilder("loadPartitions")
.addAnnotation(Override.class)
.addModifiers(Modifier.PUBLIC)
.returns(partitionsClass);

builder.addStatement("return PARTITIONS.getValue()");

return builder.build();
}

private FieldSpec partitionDataField() {
FieldSpec.Builder builder = FieldSpec.builder(String.class, "DEFAULT_PARTITION_DATA", Modifier.PRIVATE,
Modifier.STATIC, Modifier.FINAL);
builder.initializer("$S", readPartitionsJson());
return builder.build();
}

private FieldSpec partitionsLazyField() {
ParameterizedTypeName lazyType = ParameterizedTypeName.get(ClassName.get(Lazy.class),
partitionsClass);
FieldSpec.Builder builder = FieldSpec.builder(lazyType, "PARTITIONS")
.addModifiers(Modifier.PRIVATE, Modifier.STATIC, Modifier.FINAL);
CodeBlock init = CodeBlock.builder()
.addStatement("new $T<>($T::doLoadPartitions)", Lazy.class, className())
.build();

builder.initializer(init);
.addAnnotation(Override.class)
.addModifiers(Modifier.PUBLIC)
.returns(partitionsClass);

builder.addStatement("return LazyPartitionsContainer.PARTITIONS");
return builder.build();
}

private MethodSpec doLoadPartitionsMethod() {
MethodSpec.Builder builder = MethodSpec.methodBuilder("doLoadPartitions")
.addModifiers(Modifier.PRIVATE, Modifier.STATIC)
.returns(partitionsClass);

builder.addStatement("return $T.fromNode($T.parser().parse(DEFAULT_PARTITION_DATA))", partitionsClass, JsonNode.class);

return builder.build();
private TypeSpec lazyPartitionsContainer() {
CodeBlock.Builder builder = CodeBlock.builder();
JsonNode node = JsonNode.parser().parse(readPartitionsJson());
codegenPartitions(builder, node);
return TypeSpec.classBuilder("LazyPartitionsContainer")
.addModifiers(Modifier.STATIC)
.addField(FieldSpec.builder(partitionsClass, "PARTITIONS", Modifier.STATIC, Modifier.FINAL)
.initializer(builder.build())
.build())
.build();
}

private String readPartitionsJson() {
Expand All @@ -132,4 +128,104 @@ private InputStream loadResource(String name) {
Validate.notNull(resourceAsStream, "Failed to load resource from %s", name);
return resourceAsStream;
}

private void codegenPartitions(CodeBlock.Builder builder, JsonNode node) {
builder.add("$T.builder()", partitionsClass);
Map<String, JsonNode> objNode = node.asObject();

JsonNode version = objNode.get(VERSION);
if (version != null) {
builder.add(".version(");
builder.add("$S", version.asString());
builder.add(")");
}

JsonNode partitions = objNode.get(PARTITIONS);
if (partitions != null) {
partitions.asArray().forEach(partNode -> {
builder.add(".addPartition(");
codegenPartition(builder, partNode);
builder.add(")");
});
}
builder.add(".build()");
}

private void codegenPartition(CodeBlock.Builder builder, JsonNode node) {
builder.add("$T.builder()", partitionClass);
Map<String, JsonNode> objNode = node.asObject();

JsonNode id = objNode.get(ID);
if (id != null) {
builder.add(".id(");
builder.add("$S", id.asString());
builder.add(")");
}

JsonNode regionRegex = objNode.get(REGION_REGEX);
if (regionRegex != null) {
builder.add(".regionRegex(");
builder.add("$S", regionRegex.asString());
builder.add(")");
}

JsonNode regions = objNode.get(REGIONS);
if (regions != null) {
// At the moment `RegionOverride.fromNode` does nothing. We need to fix it here **and** if we keep the
// loading from textual JSON also fix `RegionOverride.fromNode`.
Map<String, JsonNode> regionsObj = regions.asObject();
regionsObj.forEach((k, v) -> {
builder.add(".putRegion($S, ", k);
codegenRegionOverride(builder, v);
builder.add(")");
});
}

JsonNode outputs = objNode.get(OUTPUTS);
if (outputs != null) {
builder.add(".outputs(");
codegenOutputs(builder, outputs);
builder.add(")");
}
builder.add(".build()");
}

private void codegenRegionOverride(CodeBlock.Builder builder, JsonNode node) {
builder.add("$T.builder().build()", regionOverrideClass);
}

private void codegenOutputs(CodeBlock.Builder builder, JsonNode node) {
builder.add("$T.builder()", outputsClass);
Map<String, JsonNode> objNode = node.asObject();

JsonNode dnsSuffix = objNode.get(DNS_SUFFIX);
if (dnsSuffix != null) {
builder.add(".dnsSuffix(");
builder.add("$S", dnsSuffix.asString());
builder.add(")");
}

JsonNode dualStackDnsSuffix = objNode.get(DUAL_STACK_DNS_SUFFIX);
if (dualStackDnsSuffix != null) {
builder.add(".dualStackDnsSuffix(");
builder.add("$S", dualStackDnsSuffix.asString());
builder.add(")");
}

JsonNode supportsFips = objNode.get(SUPPORTS_FIPS);
if (supportsFips != null) {
builder.add(".supportsFips(");
builder.add("$L", supportsFips.asBoolean());
builder.add(")");
}

JsonNode supportsDualStack = objNode.get(SUPPORTS_DUAL_STACK);
if (supportsDualStack != null) {
builder.add(".supportsDualStack(");
builder.add("$L", supportsDualStack.asBoolean());
builder.add(")");
}
builder.add(".build()");
}
}