Skip to content

Commit 970405e

Browse files
adwsinghrhernandez35
authored andcommitted
Generate SchemaIndex to support fetching Schema's from ShapeId
Fix schemaindex
1 parent 1111a2b commit 970405e

File tree

18 files changed

+322
-17
lines changed

18 files changed

+322
-17
lines changed

client/client-mock-plugin/src/main/java/software/amazon/smithy/java/client/http/mock/MockService.java

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

88
import java.util.List;
99
import software.amazon.smithy.java.core.schema.Schema;
10+
import software.amazon.smithy.java.core.schema.SchemaIndex;
1011
import software.amazon.smithy.java.core.schema.SerializableStruct;
1112
import software.amazon.smithy.java.core.serde.TypeRegistry;
1213
import software.amazon.smithy.java.server.Operation;
@@ -39,4 +40,9 @@ public Schema schema() {
3940
public TypeRegistry typeRegistry() {
4041
return TypeRegistry.empty();
4142
}
43+
44+
@Override
45+
public SchemaIndex schemaIndex() {
46+
throw new UnsupportedOperationException();
47+
}
4248
}
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+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package software.amazon.smithy.java.codegen.generators;
7+
8+
import java.util.HashMap;
9+
import java.util.Map;
10+
import java.util.function.Consumer;
11+
import software.amazon.smithy.codegen.core.directed.CustomizeDirective;
12+
import software.amazon.smithy.java.codegen.CodeGenerationContext;
13+
import software.amazon.smithy.java.codegen.CodegenUtils;
14+
import software.amazon.smithy.java.codegen.JavaCodegenSettings;
15+
import software.amazon.smithy.java.codegen.writer.JavaWriter;
16+
import software.amazon.smithy.java.core.schema.Schema;
17+
import software.amazon.smithy.java.core.schema.SchemaIndex;
18+
import software.amazon.smithy.model.shapes.ShapeId;
19+
import software.amazon.smithy.utils.SmithyInternalApi;
20+
21+
@SmithyInternalApi
22+
public final class SchemaIndexGenerator
23+
implements Consumer<CustomizeDirective<CodeGenerationContext, JavaCodegenSettings>> {
24+
25+
@Override
26+
public void accept(CustomizeDirective<CodeGenerationContext, JavaCodegenSettings> directive) {
27+
var className = "GeneratedSchemaIndex";
28+
var fileName = String
29+
.format("./%s/model/%s.java", directive.settings().packageNamespace().replace(".", "/"), className);
30+
31+
directive.context()
32+
.writerDelegator()
33+
.useFileWriter(fileName,
34+
CodegenUtils.getModelNamespace(directive.settings()),
35+
writer -> generateSchemaIndexClass(writer, className, directive));
36+
37+
// Generate META-INF/services file
38+
var serviceFileName = "./META-INF/services/" + SchemaIndex.class.getName();
39+
var schemaIndexClassName = CodegenUtils.getModelNamespace(directive.settings()) + "." + className;
40+
41+
directive.context()
42+
.writerDelegator()
43+
.useFileWriter(serviceFileName, writer -> writer.write(schemaIndexClassName));
44+
}
45+
46+
private void generateSchemaIndexClass(
47+
JavaWriter writer,
48+
String className,
49+
CustomizeDirective<CodeGenerationContext, JavaCodegenSettings> directive
50+
) {
51+
52+
writer.putContext("schemaIndex", SchemaIndex.class);
53+
writer.putContext("schema", Schema.class);
54+
writer.putContext("shapeId", ShapeId.class);
55+
writer.putContext("map", Map.class);
56+
writer.putContext("hashMap", HashMap.class);
57+
58+
var template = """
59+
/**
60+
* Generated SchemaIndex implementation that provides access to all schemas in the model.
61+
*/
62+
public final class ${className:L} extends ${schemaIndex:T} {
63+
64+
private static final ${map:T}<${shapeId:T}, ${schema:T}> SCHEMA_MAP = new ${hashMap:T}<>();
65+
66+
static {
67+
${schemaInitializers:C|}
68+
}
69+
70+
@Override
71+
public ${schema:T} getSchema(${shapeId:T} id) {
72+
return SCHEMA_MAP.get(id);
73+
}
74+
}
75+
""";
76+
77+
writer.pushState();
78+
writer.putContext("className", className);
79+
writer.putContext("schemaIndex", SchemaIndex.class);
80+
writer.putContext("schema", Schema.class);
81+
writer.putContext("shapeId", ShapeId.class);
82+
writer.putContext("schemaInitializers", new SchemaInitializersGenerator(writer, directive));
83+
writer.write(template);
84+
writer.popState();
85+
}
86+
87+
private record SchemaInitializersGenerator(
88+
JavaWriter writer,
89+
CustomizeDirective<CodeGenerationContext, JavaCodegenSettings> directive)
90+
implements Runnable {
91+
92+
@Override
93+
public void run() {
94+
var order = directive.context().schemaFieldOrder();
95+
96+
for (var shapeOrder : order.partitions()) {
97+
for (var schemaField : shapeOrder) {
98+
if (schemaField.isExternal()) {
99+
continue;
100+
}
101+
102+
var schemaReference = schemaField.className() + "." + schemaField.fieldName();
103+
104+
writer.pushState();
105+
writer.putContext("schemaReference", schemaReference);
106+
writer.write("SCHEMA_MAP.put(${schemaReference:L}.id(), ${schemaReference:L});");
107+
writer.popState();
108+
}
109+
}
110+
}
111+
}
112+
}

codegen/plugins/client-codegen/src/main/java/software/amazon/smithy/java/codegen/client/DirectedJavaClientCodegen.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import software.amazon.smithy.java.codegen.generators.MapGenerator;
1919
import software.amazon.smithy.java.codegen.generators.OperationGenerator;
2020
import software.amazon.smithy.java.codegen.generators.ResourceGenerator;
21+
import software.amazon.smithy.java.codegen.generators.SchemaIndexGenerator;
2122
import software.amazon.smithy.java.codegen.generators.SchemasGenerator;
2223
import software.amazon.smithy.java.codegen.generators.ServiceExceptionGenerator;
2324
import software.amazon.smithy.java.codegen.generators.SharedSerdeGenerator;
@@ -128,6 +129,7 @@ public void customizeBeforeIntegrations(CustomizeDirective<CodeGenerationContext
128129
if (!directive.settings().useExternalTypes()) {
129130
new SchemasGenerator().accept(directive);
130131
new SharedSerdeGenerator().accept(directive);
132+
new SchemaIndexGenerator().accept(directive);
131133
}
132134
}
133135
}

codegen/plugins/server-codegen/src/main/java/software/amazon/smithy/java/codegen/server/DirectedJavaServerCodegen.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
import software.amazon.smithy.java.codegen.generators.MapGenerator;
3030
import software.amazon.smithy.java.codegen.generators.OperationGenerator;
3131
import software.amazon.smithy.java.codegen.generators.ResourceGenerator;
32+
import software.amazon.smithy.java.codegen.generators.SchemaIndexGenerator;
3233
import software.amazon.smithy.java.codegen.generators.SchemasGenerator;
3334
import software.amazon.smithy.java.codegen.generators.ServiceExceptionGenerator;
3435
import software.amazon.smithy.java.codegen.generators.SharedSerdeGenerator;
@@ -139,6 +140,7 @@ public void customizeBeforeIntegrations(CustomizeDirective<CodeGenerationContext
139140
if (!directive.settings().useExternalTypes()) {
140141
new SharedSerdeGenerator().accept(directive);
141142
new SchemasGenerator().accept(directive);
143+
new SchemaIndexGenerator().accept(directive);
142144
}
143145
}
144146
}

codegen/plugins/server-codegen/src/main/java/software/amazon/smithy/java/codegen/server/ServerSymbolProperties.java

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,4 +36,9 @@ private ServerSymbolProperties() {}
3636
* Name to use for the operation when used as a field inside the service.
3737
*/
3838
public static final Property<String> OPERATION_FIELD_NAME = Property.named("operation-field-name");
39+
40+
/**
41+
* Namespace under which types of this service would be generated.
42+
*/
43+
public static final Property<String> TYPES_NAMESPACE = Property.named("types-namespace");
3944
}

codegen/plugins/server-codegen/src/main/java/software/amazon/smithy/java/codegen/server/ServiceJavaSymbolProvider.java

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,7 @@ private Symbol getServerJavaClassSymbol() {
7070
CodegenUtils.getServiceExceptionSymbol(packageNamespace(), serviceName))
7171
.putProperty(SymbolProperties.SERVICE_API_SERVICE,
7272
CodegenUtils.getServiceApiSymbol(packageNamespace(), serviceName))
73+
.putProperty(ServerSymbolProperties.TYPES_NAMESPACE, format("%s.model", packageNamespace()))
7374
.namespace(format("%s.service", packageNamespace()), ".")
7475
.declarationFile(format("./%s/service/%s.java", packageNamespace().replace(".", "/"), serviceName))
7576
.build();

codegen/plugins/server-codegen/src/main/java/software/amazon/smithy/java/codegen/server/generators/ServiceGenerator.java

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import software.amazon.smithy.java.codegen.server.ServerSymbolProperties;
2424
import software.amazon.smithy.java.codegen.writer.JavaWriter;
2525
import software.amazon.smithy.java.core.schema.Schema;
26+
import software.amazon.smithy.java.core.schema.SchemaIndex;
2627
import software.amazon.smithy.java.core.schema.SerializableStruct;
2728
import software.amazon.smithy.java.core.serde.TypeRegistry;
2829
import software.amazon.smithy.java.framework.model.UnknownOperationException;
@@ -55,6 +56,8 @@ public void accept(
5556
})
5657
.toList();
5758
var operations = operationsInfo.stream().map(OperationInfo::symbol).toList();
59+
var generatedSchemaIndex =
60+
directive.symbol().expectProperty(ServerSymbolProperties.TYPES_NAMESPACE) + ".GeneratedSchemaIndex";
5861
directive.context().writerDelegator().useShapeWriter(shape, writer -> {
5962
writer.pushState(new ClassSection(shape));
6063
var template =
@@ -73,6 +76,8 @@ public final class ${service:T} implements ${serviceType:T} {
7376
7477
${builder:C|}
7578
79+
private static final ${schemaIndex:T} SCHEMA_INDEX = new ${generatedSchemaIndex:L}();
80+
7681
@Override
7782
@SuppressWarnings("unchecked")
7883
public <I extends ${serializableStruct:T}, O extends ${serializableStruct:T}> ${operationHolder:T}<I, O> getOperation(String operationName) {
@@ -93,6 +98,11 @@ public final class ${service:T} implements ${serviceType:T} {
9398
public ${typeRegistryClass:T} typeRegistry() {
9499
return TYPE_REGISTRY;
95100
}
101+
102+
@Override
103+
public ${schemaIndex:T} schemaIndex() {
104+
return SCHEMA_INDEX;
105+
}
96106
}
97107
""";
98108
writer.putContext("operationHolder", Operation.class);
@@ -102,6 +112,8 @@ public final class ${service:T} implements ${serviceType:T} {
102112
writer.putContext("service", directive.symbol());
103113
writer.putContext("id", new IdStringGenerator(writer, shape));
104114
writer.putContext("typeRegistryClass", TypeRegistry.class);
115+
writer.putContext("schemaIndex", SchemaIndex.class);
116+
writer.putContext("generatedSchemaIndex", generatedSchemaIndex);
105117
var errorSymbols = getImplicitErrorSymbols(
106118
directive.symbolProvider(),
107119
directive.model(),

codegen/plugins/types-codegen/src/main/java/software/amazon/smithy/java/codegen/types/DirectedJavaTypeCodegen.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import software.amazon.smithy.java.codegen.generators.EnumGenerator;
2626
import software.amazon.smithy.java.codegen.generators.ListGenerator;
2727
import software.amazon.smithy.java.codegen.generators.MapGenerator;
28+
import software.amazon.smithy.java.codegen.generators.SchemaIndexGenerator;
2829
import software.amazon.smithy.java.codegen.generators.SchemasGenerator;
2930
import software.amazon.smithy.java.codegen.generators.SharedSerdeGenerator;
3031
import software.amazon.smithy.java.codegen.generators.StructureGenerator;
@@ -101,6 +102,7 @@ public void generateIntEnumShape(GenerateIntEnumDirective<CodeGenerationContext,
101102
public void customizeBeforeIntegrations(CustomizeDirective<CodeGenerationContext, JavaCodegenSettings> directive) {
102103
new SharedSerdeGenerator().accept(directive);
103104
new SchemasGenerator().accept(directive);
105+
new SchemaIndexGenerator().accept(directive);
104106
}
105107

106108
@Override

codegen/plugins/types-codegen/src/test/java/software/amazon/smithy/java/codegen/types/CodegenTest.java

Lines changed: 20 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,7 @@
55

66
package software.amazon.smithy.java.codegen.types;
77

8-
import static org.hamcrest.MatcherAssert.assertThat;
9-
import static org.hamcrest.Matchers.containsInAnyOrder;
8+
import static org.assertj.core.api.Assertions.assertThat;
109
import static org.junit.jupiter.api.Assertions.assertEquals;
1110

1211
import java.net.URL;
@@ -48,17 +47,18 @@ void expectedFilesExist() {
4847
var settings = settingsBuilder.build();
4948
var context = contextBuilder.settings(settings).build();
5049
plugin.execute(context);
51-
assertEquals(7, manifest.getFiles().size());
52-
assertThat(
53-
manifest.getFiles(),
54-
containsInAnyOrder(
50+
assertThat(manifest.getFiles())
51+
.hasSize(9)
52+
.containsExactlyInAnyOrder(
5553
Path.of("/test/smithy/codegen/types/test/model/EnumShape.java"),
5654
Path.of("/test/smithy/codegen/types/test/model/IntEnumShape.java"),
5755
Path.of("/test/smithy/codegen/types/test/model/Schemas.java"),
5856
Path.of("/test/smithy/codegen/types/test/model/SharedSerde.java"),
5957
Path.of("/test/smithy/codegen/types/test/model/StructureShape.java"),
6058
Path.of("/test/smithy/codegen/types/test/model/UnionShape.java"),
61-
Path.of("/META-INF/smithy-java/type-mappings.properties")));
59+
Path.of("/test/smithy/codegen/types/test/model/GeneratedSchemaIndex.java"),
60+
Path.of("/META-INF/smithy-java/type-mappings.properties"),
61+
Path.of("/META-INF/services/software.amazon.smithy.java.core.schema.SchemaIndex"));
6262
}
6363

6464
@Test
@@ -68,14 +68,15 @@ void respectsSelector() {
6868
.build();
6969
var context = contextBuilder.settings(settings).build();
7070
plugin.execute(context);
71-
assertEquals(4, manifest.getFiles().size());
72-
assertThat(
73-
manifest.getFiles(),
74-
containsInAnyOrder(
71+
assertThat(manifest.getFiles())
72+
.hasSize(6)
73+
.containsExactlyInAnyOrder(
7574
Path.of("/test/smithy/codegen/types/test/model/Schemas.java"),
7675
Path.of("/test/smithy/codegen/types/test/model/SharedSerde.java"),
7776
Path.of("/test/smithy/codegen/types/test/model/StructureShape.java"),
78-
Path.of("/META-INF/smithy-java/type-mappings.properties")));
77+
Path.of("/test/smithy/codegen/types/test/model/GeneratedSchemaIndex.java"),
78+
Path.of("/META-INF/smithy-java/type-mappings.properties"),
79+
Path.of("/META-INF/services/software.amazon.smithy.java.core.schema.SchemaIndex"));
7980
}
8081

8182
@Test
@@ -86,15 +87,17 @@ void specificShapesAdded() {
8687
.build();
8788
var context = contextBuilder.settings(settings).build();
8889
plugin.execute(context);
89-
assertEquals(5, manifest.getFiles().size());
90-
assertThat(
91-
manifest.getFiles(),
92-
containsInAnyOrder(
90+
assertEquals(7, manifest.getFiles().size());
91+
assertThat(manifest.getFiles())
92+
.hasSize(7)
93+
.containsExactlyInAnyOrder(
9394
Path.of("/test/smithy/codegen/types/test/model/Schemas.java"),
9495
Path.of("/test/smithy/codegen/types/test/model/SharedSerde.java"),
9596
Path.of("/test/smithy/codegen/types/test/model/StructureShape.java"),
9697
Path.of("/test/smithy/codegen/types/test/model/UnionShape.java"),
97-
Path.of("/META-INF/smithy-java/type-mappings.properties")));
98+
Path.of("/test/smithy/codegen/types/test/model/GeneratedSchemaIndex.java"),
99+
Path.of("/META-INF/smithy-java/type-mappings.properties"),
100+
Path.of("/META-INF/services/software.amazon.smithy.java.core.schema.SchemaIndex"));
98101
}
99102

100103
}
Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package software.amazon.smithy.java.core.schema;
7+
8+
import java.util.List;
9+
import java.util.ServiceLoader;
10+
import software.amazon.smithy.model.shapes.ShapeId;
11+
12+
public abstract class SchemaIndex {
13+
14+
private static final SchemaIndex COMBINED_SCHEMA_INDEX = createCombinedSchemaIndex();
15+
16+
private static SchemaIndex createCombinedSchemaIndex() {
17+
return new CombinedSchemaIndex(
18+
ServiceLoader.load(SchemaIndex.class).stream().map(ServiceLoader.Provider::get).toList());
19+
}
20+
21+
public static SchemaIndex getCombinedSchemaIndex() {
22+
return COMBINED_SCHEMA_INDEX;
23+
}
24+
25+
public abstract Schema getSchema(ShapeId id);
26+
27+
private static final class CombinedSchemaIndex extends SchemaIndex {
28+
29+
private final List<SchemaIndex> indexes;
30+
31+
public CombinedSchemaIndex(List<SchemaIndex> indexes) {
32+
this.indexes = indexes;
33+
}
34+
35+
@Override
36+
public Schema getSchema(ShapeId id) {
37+
for (var index : indexes) {
38+
var schema = index.getSchema(id);
39+
if (schema != null) {
40+
return schema;
41+
}
42+
}
43+
throw new IllegalArgumentException("No schema found for id `" + id + "`");
44+
}
45+
}
46+
}

0 commit comments

Comments
 (0)