fix(sdk-codegen): update errorType in waitable traits to use correct …

…shape Id instead of awsquery code
aws · Sep 21, 2024 · aa72922 · aa72922
1 parent 70db983
commit aa72922
Show file tree

Hide file tree

Showing 2 changed files with 90 additions and 0 deletions.
diff --git a/...n/src/main/java/software/amazon/smithy/aws/typescript/codegen/ProcessAwsQueryWaiters.java b/...n/src/main/java/software/amazon/smithy/aws/typescript/codegen/ProcessAwsQueryWaiters.java
@@ -0,0 +1,89 @@
+package software.amazon.smithy.aws.typescript.codegen;
+
+import software.amazon.smithy.aws.traits.protocols.AwsQueryErrorTrait;
+import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.node.ObjectNode;
+import software.amazon.smithy.model.shapes.*;
+import software.amazon.smithy.model.traits.ErrorTrait;
+import software.amazon.smithy.model.transform.ModelTransformer;
+import software.amazon.smithy.typescript.codegen.TypeScriptSettings;
+import software.amazon.smithy.typescript.codegen.integration.TypeScriptIntegration;
+import software.amazon.smithy.utils.SmithyInternalApi;
+import software.amazon.smithy.model.knowledge.TopDownIndex;
+import software.amazon.smithy.waiters.Acceptor;
+import software.amazon.smithy.waiters.Matcher;
+import software.amazon.smithy.waiters.WaitableTrait;
+import software.amazon.smithy.waiters.Waiter;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.ArrayList;
+import java.util.List;
+
+
+@SmithyInternalApi
+public final class ProcessAwsQueryWaiters implements TypeScriptIntegration {
+
+    @Override
+    public Model preprocessModel(Model model, TypeScriptSettings settings) {
+        Map<String, String> errorCodeToShapeId = new HashMap<>();
+
+        ServiceShape serviceShape = settings.getService(model);
+        TopDownIndex topDownIndex = TopDownIndex.of(model);
+        for (OperationShape operationShape : topDownIndex.getContainedOperations(serviceShape)) {
+            for (ShapeId errorShapeId : operationShape.getErrors()) {
+                Shape errorShape = model.expectShape(errorShapeId);
+                if (errorShape.hasTrait(ErrorTrait.class) && errorShape.hasTrait(AwsQueryErrorTrait.class)) {
+                    AwsQueryErrorTrait awsQueryTrait = errorShape.expectTrait(AwsQueryErrorTrait.class);
+                    errorCodeToShapeId.put(awsQueryTrait.getCode(), errorShape.getId().getName());
+                }
+            }
+        }
+
+        List<Shape> modifiedShapes = new ArrayList<>();
+
+        for (OperationShape operationShape : topDownIndex.getContainedOperations(serviceShape)) {
+            OperationShape.Builder operationBuilder = operationShape.toBuilder();
+            if (operationShape.hasTrait(WaitableTrait.class)) {
+                operationBuilder.removeTrait(WaitableTrait.ID);
+                WaitableTrait waiterTrait = operationShape.expectTrait(WaitableTrait.class);
+                WaitableTrait.Builder waitableTraitBuilder = (WaitableTrait.Builder) waiterTrait.toBuilder();
+                for (Map.Entry<String, Waiter> entry : waiterTrait.getWaiters().entrySet()){
+                    String name = entry.getKey();
+                    Waiter waiter = entry.getValue();
+                    Waiter.Builder waiterBuilder = (Waiter.Builder) waiter.toBuilder();
+                    waiterBuilder.clearAcceptors();
+                    for (Acceptor acceptor : waiter.getAcceptors()){
+                        ObjectNode acceptorNode = acceptor.toNode().expectObjectNode();
+                        Matcher matcher = acceptor.getMatcher();
+                        if (matcher instanceof Matcher.ErrorTypeMember){
+                            ObjectNode matcherNode = matcher.toNode().expectObjectNode();
+
+                            String errorCode = matcherNode.expectStringMember("errorType").getValue();
+                            if (errorCodeToShapeId.containsKey(errorCode)){
+                                matcherNode = matcherNode.toBuilder()
+                                        .withMember("errorType", errorCodeToShapeId.get(errorCode))
+                                        .build();
+
+                                acceptorNode = acceptorNode.toBuilder()
+                                        .withMember("matcher", matcherNode)
+                                        .build();
+                            }
+                        }
+                        waiterBuilder.addAcceptor(Acceptor.fromNode(acceptorNode));
+                    }
+                    waitableTraitBuilder.put(name,waiterBuilder.build());
+                }
+                operationBuilder.addTrait(waitableTraitBuilder.build());
+            }
+            modifiedShapes.add(operationBuilder.build());
+        }
+
+        if (!modifiedShapes.isEmpty()) {
+            ModelTransformer transformer = ModelTransformer.create();
+            model = transformer.replaceShapes(model, modifiedShapes);
+        }
+
+        return model;
+    }
+}
diff --git a/...-INF/services/software.amazon.smithy.typescript.codegen.integration.TypeScriptIntegration b/...-INF/services/software.amazon.smithy.typescript.codegen.integration.TypeScriptIntegration
@@ -27,6 +27,7 @@ software.amazon.smithy.aws.typescript.codegen.AddDocumentClientPlugin
 software.amazon.smithy.aws.typescript.codegen.AddEndpointDiscoveryPlugin
 software.amazon.smithy.aws.typescript.codegen.AddHttpChecksumDependency
 software.amazon.smithy.aws.typescript.codegen.AddSigv4aPlugin
+software.amazon.smithy.aws.typescript.codegen.ProcessAwsQueryWaiters
 software.amazon.smithy.aws.typescript.codegen.auth.http.integration.AwsSdkCustomizeHttpBearerTokenAuth
 software.amazon.smithy.aws.typescript.codegen.auth.http.integration.SupportSigV4Auth
 software.amazon.smithy.aws.typescript.codegen.auth.http.integration.AwsSdkCustomizeSigV4Auth