Skip to content

Commit aa57647

Browse files
author
Ran Vaknin
committed
fix(sdk-codegen): update errorType in waitable traits to use correct error codes
1 parent 9ed53c0 commit aa57647

File tree

2 files changed

+145
-0
lines changed

2 files changed

+145
-0
lines changed
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,144 @@
1+
package software.amazon.smithy.aws.typescript.codegen;
2+
3+
import software.amazon.smithy.model.Model;
4+
import software.amazon.smithy.model.node.ArrayNode;
5+
import software.amazon.smithy.model.node.Node;
6+
import software.amazon.smithy.model.node.ObjectNode;
7+
import software.amazon.smithy.model.node.StringNode;
8+
import software.amazon.smithy.model.shapes.AbstractShapeBuilder;
9+
import software.amazon.smithy.model.shapes.OperationShape;
10+
import software.amazon.smithy.model.shapes.Shape;
11+
import software.amazon.smithy.model.shapes.ShapeId;
12+
import software.amazon.smithy.model.traits.DynamicTrait;
13+
import software.amazon.smithy.model.traits.Trait;
14+
import software.amazon.smithy.model.transform.ModelTransformer;
15+
import software.amazon.smithy.typescript.codegen.TypeScriptSettings;
16+
import software.amazon.smithy.typescript.codegen.integration.TypeScriptIntegration;
17+
import software.amazon.smithy.utils.SmithyInternalApi;
18+
19+
import java.util.*;
20+
21+
@SmithyInternalApi
22+
public final class ProcessAwsQueryWaiters implements TypeScriptIntegration {
23+
24+
@Override
25+
public Model preprocessModel(Model model, TypeScriptSettings settings) {
26+
// create mapping of exception shape IDs and awsQuery error code
27+
Map<String, ShapeId> errorCodeToShapeId = new HashMap<>();
28+
for (Shape shape : model.toSet()) {
29+
if (shape.hasTrait("smithy.api#error") && shape.hasTrait("aws.protocols#awsQueryError")) {
30+
Optional<Trait> awsQueryTraitOpt = shape.findTrait("aws.protocols#awsQueryError");
31+
if (awsQueryTraitOpt.isPresent()) {
32+
Trait awsQueryTrait = awsQueryTraitOpt.get();
33+
ObjectNode traitValue = awsQueryTrait.toNode().expectObjectNode();
34+
Optional<Node> memberNodeOpt = traitValue.getMember("code");
35+
if (memberNodeOpt.isPresent()){
36+
Optional<StringNode> codeNodeOpt = memberNodeOpt.get().asStringNode();
37+
if (codeNodeOpt.isPresent()){
38+
String code = codeNodeOpt.get().getValue();
39+
errorCodeToShapeId.put(code, shape.getId());
40+
}
41+
}
42+
}
43+
}
44+
}
45+
46+
47+
List<Shape> modifiedShapes = new ArrayList<>();
48+
49+
for (Shape shape : model.toSet()) {
50+
Optional<Trait> waitableTraitOpt = shape.findTrait("smithy.waiters#waitable");
51+
if (waitableTraitOpt.isPresent()) {
52+
Trait waitableTrait = waitableTraitOpt.get();
53+
ObjectNode traitValue = waitableTrait.toNode().expectObjectNode();
54+
ObjectNode.Builder traitValueBuilder = traitValue.toBuilder();
55+
boolean modified = false;
56+
57+
for (Map.Entry<StringNode, Node> waiterEntry : traitValue.getMembers().entrySet()) {
58+
StringNode waiterNameNode = waiterEntry.getKey();
59+
Node waiterNode = waiterEntry.getValue();
60+
61+
if (waiterNode.isObjectNode()) {
62+
ObjectNode waiterObject = waiterNode.expectObjectNode();
63+
ObjectNode.Builder waiterObjectBuilder = waiterObject.toBuilder();
64+
boolean waiterModified = false;
65+
66+
Optional<ArrayNode> acceptorsArrayOpt = waiterObject.getArrayMember("acceptors");
67+
if (acceptorsArrayOpt.isPresent()) {
68+
ArrayNode acceptorsArray = acceptorsArrayOpt.get();
69+
ArrayNode.Builder acceptorsArrayBuilder = ArrayNode.builder();
70+
boolean acceptorsModified = false;
71+
72+
for (Node acceptorNode : acceptorsArray.getElements()) {
73+
if (acceptorNode.isObjectNode()) {
74+
ObjectNode acceptorObject = acceptorNode.expectObjectNode();
75+
Optional<ObjectNode> matcherObjectOpt = acceptorObject.getObjectMember("matcher");
76+
if (matcherObjectOpt.isPresent()) {
77+
ObjectNode matcherObject = matcherObjectOpt.get();
78+
Optional<StringNode> errorTypeNodeOpt = matcherObject.getStringMember("errorType");
79+
if (errorTypeNodeOpt.isPresent()) {
80+
String errorType = errorTypeNodeOpt.get().getValue();
81+
ShapeId errorShapeId = errorCodeToShapeId.get(errorType);
82+
if (errorShapeId != null) {
83+
// Replace the errorType value with the shape ID
84+
ObjectNode modifiedMatcherObject = matcherObject.withMember("errorType", Node.from(errorShapeId.toString()));
85+
ObjectNode modifiedAcceptorObject = acceptorObject.withMember("matcher", modifiedMatcherObject);
86+
acceptorsArrayBuilder.withValue(modifiedAcceptorObject);
87+
acceptorsModified = true;
88+
continue;
89+
}
90+
}
91+
}
92+
// If not modified, add the original acceptor
93+
acceptorsArrayBuilder.withValue(acceptorObject);
94+
} else {
95+
acceptorsArrayBuilder.withValue(acceptorNode);
96+
}
97+
}
98+
if (acceptorsModified) {
99+
waiterObjectBuilder.withMember("acceptors", acceptorsArrayBuilder.build());
100+
waiterModified = true;
101+
}
102+
}
103+
if (waiterModified) {
104+
traitValueBuilder.withMember(waiterNameNode, waiterObjectBuilder.build());
105+
modified = true;
106+
}
107+
}
108+
}
109+
110+
if (modified) {
111+
Trait modifiedWaitableTrait = new DynamicTrait(waitableTrait.toShapeId(), traitValueBuilder.build());
112+
AbstractShapeBuilder<?, ?> shapeBuilder = Shape.shapeToBuilder(shape);
113+
114+
shapeBuilder.removeTrait(waitableTrait.toShapeId());
115+
shapeBuilder.addTrait(modifiedWaitableTrait);
116+
Shape modifiedShape = shapeBuilder.build();
117+
118+
modifiedShapes.add(modifiedShape);
119+
}
120+
}
121+
}
122+
123+
// replace the modified shapes in the model
124+
if (!modifiedShapes.isEmpty()) {
125+
ModelTransformer transformer = ModelTransformer.create();
126+
model = transformer.replaceShapes(model, modifiedShapes);
127+
}
128+
129+
// print the transformation
130+
for (Shape updatedShape : model.toSet()) {
131+
Optional<Trait> waitableTraitOpt = updatedShape.findTrait("smithy.waiters#waitable");
132+
if (waitableTraitOpt.isPresent()) {
133+
Trait waitableTrait = waitableTraitOpt.get();
134+
ObjectNode traitValue = waitableTrait.toNode().expectObjectNode();
135+
136+
System.out.println("Waitable trait for shape: " + updatedShape.getId());
137+
String json = Node.prettyPrintJson(traitValue);
138+
System.out.println(json);
139+
}
140+
}
141+
142+
return model;
143+
}
144+
}

codegen/smithy-aws-typescript-codegen/src/main/resources/META-INF/services/software.amazon.smithy.typescript.codegen.integration.TypeScriptIntegration

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,3 +33,4 @@ software.amazon.smithy.aws.typescript.codegen.auth.http.integration.AwsSdkCustom
3333
software.amazon.smithy.aws.typescript.codegen.auth.http.integration.AddAwsDefaultSigningName
3434
software.amazon.smithy.aws.typescript.codegen.auth.http.integration.AddSTSAuthCustomizations
3535
software.amazon.smithy.aws.typescript.codegen.auth.http.integration.AwsSdkCustomizeEndpointRuleSetHttpAuthSchemeProvider
36+
software.amazon.smithy.aws.typescript.codegen.ProcessAwsQueryWaiters

0 commit comments

Comments
 (0)