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
+ }
0 commit comments