Add support for retryable trait

codegen/src/main/java/software/amazon/awssdk/codegen/AddExceptionShapes.java

@@ -44,9 +44,10 @@ private Map<String, ShapeModel> constructExceptionShapes() {
         // Java shape models, to be constructed
         Map<String, ShapeModel> javaShapes = new HashMap<>();
 
-        for (Map.Entry<String, Shape> shape : getServiceModel().getShapes().entrySet()) {
-            if (shape.getValue().isException()) {
-                String errorShapeName = shape.getKey();
+        for (Map.Entry<String, Shape> kvp : getServiceModel().getShapes().entrySet()) {
+            if (kvp.getValue().isException()) {
+                Shape shape = kvp.getValue();
+                String errorShapeName = kvp.getKey();
                 String javaClassName = getNamingStrategy().getExceptionName(errorShapeName);
 
                 ShapeModel exceptionShapeModel = generateShapeModel(javaClassName,
@@ -55,8 +56,10 @@ private Map<String, ShapeModel> constructExceptionShapes() {
                 exceptionShapeModel.setType(ShapeType.Exception.getValue());
                 exceptionShapeModel.setErrorCode(getErrorCode(errorShapeName));
                 exceptionShapeModel.setHttpStatusCode(getHttpStatusCode(errorShapeName));
+                exceptionShapeModel.withIsRetryable(shape.isRetryable());
+                exceptionShapeModel.withIsThrottling(shape.isThrottling());
                 if (exceptionShapeModel.getDocumentation() == null) {
-                    exceptionShapeModel.setDocumentation(shape.getValue().getDocumentation());
+                    exceptionShapeModel.setDocumentation(shape.getDocumentation());
                 }
 
                 javaShapes.put(javaClassName, exceptionShapeModel);

codegen/src/main/java/software/amazon/awssdk/codegen/AddShapes.java

@@ -87,6 +87,8 @@ protected final ShapeModel generateShapeModel(String javaClassName, String shape
         shapeModel.withXmlNamespace(shape.getXmlNamespace());
         shapeModel.withIsUnion(shape.isUnion());
         shapeModel.withIsFault(shape.isFault());
+        shapeModel.withIsRetryable(shape.isRetryable());
+        shapeModel.withIsThrottling(shape.isThrottling());
 
         boolean hasHeaderMember = false;
         boolean hasStatusCodeMember = false;

codegen/src/main/java/software/amazon/awssdk/codegen/model/intermediate/ShapeModel.java

@@ -75,6 +75,9 @@ public class ShapeModel extends DocumentationModel implements HasDeprecation {
 
     private boolean union;
 
+    private boolean retryable;
+    private boolean throttling;
+
     public ShapeModel() {
     }
 
@@ -648,4 +651,22 @@ public ShapeModel withIsFault(boolean fault) {
         this.fault = fault;
         return this;
     }
+
+    public boolean isRetryable() {
+        return retryable;
+    }
+
+    public ShapeModel withIsRetryable(boolean retryable) {
+        this.retryable = retryable;
+        return this;
+    }
+
+    public boolean isThrottling() {
+        return throttling;
+    }
+
+    public ShapeModel withIsThrottling(boolean throttling) {
+        this.throttling = throttling;
+        return this;
+    }
 }

codegen/src/main/java/software/amazon/awssdk/codegen/model/service/RetryableTrait.java

@@ -0,0 +1,33 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.codegen.model.service;
+
+public class RetryableTrait {
+
+    private Boolean throttling;
+
+    public void setThrottling(boolean throttling) {
+        this.throttling = throttling;
+    }
+
+    public Boolean getThrottling() {
+        return throttling;
+    }
+
+    public boolean isThrottling() {
+        return Boolean.TRUE.equals(throttling);
+    }
+}

codegen/src/main/java/software/amazon/awssdk/codegen/model/service/Shape.java

@@ -78,6 +78,8 @@ public class Shape {
 
     private boolean union;
 
+    private RetryableTrait retryable;
+
     public boolean isFault() {
         return fault;
     }
@@ -345,4 +347,16 @@ public boolean isUnion() {
     public void setUnion(boolean union) {
         this.union = union;
     }
+
+    public void setRetryable(RetryableTrait retryable) {
+        this.retryable = retryable;
+    }
+
+    public boolean isRetryable() {
+        return retryable != null;
+    }
+
+    public boolean isThrottling() {
+        return retryable != null && retryable.isThrottling();
+    }
 }

codegen/src/main/java/software/amazon/awssdk/codegen/poet/model/AwsServiceModel.java

@@ -34,6 +34,7 @@
 import com.squareup.javapoet.WildcardTypeName;
 import java.io.Serializable;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
@@ -437,6 +438,7 @@ private List<MethodSpec> modelClassMethods() {
                 methodSpecs.add(builderMethod());
                 methodSpecs.add(serializableBuilderClass());
                 methodSpecs.addAll(memberGetters());
+                methodSpecs.addAll(retryableOverrides());
                 break;
             default:
                 methodSpecs.addAll(addModifier(memberGetters(), FINAL));
@@ -689,6 +691,28 @@ private CodeBlock getterStatement(MemberModel model) {
         return CodeBlock.of("return $N;", modelVariable.getVariableName());
     }
 
+    private List<MethodSpec> retryableOverrides() {
+        if (shapeModel.isRetryable()) {
+            MethodSpec isRetryable = MethodSpec.methodBuilder("isRetryableException")
+                                               .addAnnotation(Override.class)
+                                               .addModifiers(PUBLIC)
+                                               .returns(TypeName.BOOLEAN)
+                                               .addStatement("return true")
+                                               .build();
+            if (shapeModel.isThrottling()) {
+                MethodSpec isThrottling = MethodSpec.methodBuilder("isThrottlingException")
+                                                   .addAnnotation(Override.class)
+                                                   .addModifiers(PUBLIC)
+                                                   .returns(TypeName.BOOLEAN)
+                                                   .addStatement("return true")
+                                                   .build();
+                return Arrays.asList(isRetryable, isThrottling);
+            }
+            return Arrays.asList(isRetryable);
+        }
+        return emptyList();
+    }
+
     private List<TypeSpec> nestedModelClassTypes() {
         List<TypeSpec> nestedClasses = new ArrayList<>();
         switch (shapeModel.getShapeType()) {

codegen/src/test/java/software/amazon/awssdk/codegen/poet/model/RetryableExceptionClassSpecTest.java

@@ -0,0 +1,61 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.codegen.poet.model;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static software.amazon.awssdk.codegen.poet.PoetMatchers.generatesTo;
+
+import java.io.File;
+import java.io.IOException;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import software.amazon.awssdk.codegen.C2jModels;
+import software.amazon.awssdk.codegen.IntermediateModelBuilder;
+import software.amazon.awssdk.codegen.model.config.customization.CustomizationConfig;
+import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
+import software.amazon.awssdk.codegen.model.intermediate.ShapeType;
+import software.amazon.awssdk.codegen.model.service.ServiceModel;
+import software.amazon.awssdk.codegen.poet.ClassSpec;
+import software.amazon.awssdk.codegen.utils.ModelLoaderUtils;
+
+public class RetryableExceptionClassSpecTest {
+    private static IntermediateModel intermediateModel;
+
+    @BeforeAll
+    public static void setUp() throws IOException {
+        File serviceModelFile =
+            new File(RetryableExceptionClassSpecTest.class.getResource("exceptions/service-2.json").getFile());
+        File customizationConfigFile = new File(RetryableExceptionClassSpecTest.class.getResource("exceptions/customization.config")
+                                                                                     .getFile());
+        intermediateModel = new IntermediateModelBuilder(
+            C2jModels.builder()
+                     .serviceModel(ModelLoaderUtils.loadModel(ServiceModel.class, serviceModelFile))
+                     .customizationConfig(ModelLoaderUtils.loadModel(CustomizationConfig.class, customizationConfigFile))
+                     .build())
+            .build();
+    }
+
+    @Test
+    public void serviceCodegen_withErrorsWithRetryableTrait_ShouldOverrideIsRetryableAndIsThrottling() {
+        intermediateModel.getShapes().forEach((name, shape) -> {
+            if (shape.getShapeType() == ShapeType.Exception) {
+                ClassSpec spec = new AwsServiceModel(intermediateModel, shape);
+                assertThat(spec, generatesTo( "exceptions/" + name.toLowerCase() + ".java"));
+            }
+        });
+    }
+
+}

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/model/exceptions/customization.config

@@ -0,0 +1,2 @@
+{}
+

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/model/exceptions/jsonserviceinternalservererrorexception.java

@@ -0,0 +1,127 @@
+package software.amazon.awssdk.services.json.model;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import software.amazon.awssdk.annotations.Generated;
+import software.amazon.awssdk.awscore.exception.AwsErrorDetails;
+import software.amazon.awssdk.core.SdkField;
+import software.amazon.awssdk.core.SdkPojo;
+import software.amazon.awssdk.utils.builder.CopyableBuilder;
+import software.amazon.awssdk.utils.builder.ToCopyableBuilder;
+
+/**
+ * <p>
+ * There was an internal server error.
+ * </p>
+ */
+@Generated("software.amazon.awssdk:codegen")
+public final class JsonServiceInternalServerErrorException extends JsonException implements
+        ToCopyableBuilder<JsonServiceInternalServerErrorException.Builder, JsonServiceInternalServerErrorException> {
+    private static final List<SdkField<?>> SDK_FIELDS = Collections.unmodifiableList(Arrays.asList());
+
+    private static final long serialVersionUID = 1L;
+
+    private JsonServiceInternalServerErrorException(BuilderImpl builder) {
+        super(builder);
+    }
+
+    @Override
+    public Builder toBuilder() {
+        return new BuilderImpl(this);
+    }
+
+    public static Builder builder() {
+        return new BuilderImpl();
+    }
+
+    public static Class<? extends Builder> serializableBuilderClass() {
+        return BuilderImpl.class;
+    }
+
+    @Override
+    public boolean isRetryableException() {
+        return true;
+    }
+
+    @Override
+    public final List<SdkField<?>> sdkFields() {
+        return SDK_FIELDS;
+    }
+
+    public interface Builder extends SdkPojo, CopyableBuilder<Builder, JsonServiceInternalServerErrorException>,
+            JsonException.Builder {
+        @Override
+        Builder awsErrorDetails(AwsErrorDetails awsErrorDetails);
+
+        @Override
+        Builder message(String message);
+
+        @Override
+        Builder requestId(String requestId);
+
+        @Override
+        Builder statusCode(int statusCode);
+
+        @Override
+        Builder cause(Throwable cause);
+
+        @Override
+        Builder writableStackTrace(Boolean writableStackTrace);
+    }
+
+    static final class BuilderImpl extends JsonException.BuilderImpl implements Builder {
+        private BuilderImpl() {
+        }
+
+        private BuilderImpl(JsonServiceInternalServerErrorException model) {
+            super(model);
+        }
+
+        @Override
+        public BuilderImpl awsErrorDetails(AwsErrorDetails awsErrorDetails) {
+            this.awsErrorDetails = awsErrorDetails;
+            return this;
+        }
+
+        @Override
+        public BuilderImpl message(String message) {
+            this.message = message;
+            return this;
+        }
+
+        @Override
+        public BuilderImpl requestId(String requestId) {
+            this.requestId = requestId;
+            return this;
+        }
+
+        @Override
+        public BuilderImpl statusCode(int statusCode) {
+            this.statusCode = statusCode;
+            return this;
+        }
+
+        @Override
+        public BuilderImpl cause(Throwable cause) {
+            this.cause = cause;
+            return this;
+        }
+
+        @Override
+        public BuilderImpl writableStackTrace(Boolean writableStackTrace) {
+            this.writableStackTrace = writableStackTrace;
+            return this;
+        }
+
+        @Override
+        public JsonServiceInternalServerErrorException build() {
+            return new JsonServiceInternalServerErrorException(this);
+        }
+
+        @Override
+        public List<SdkField<?>> sdkFields() {
+            return SDK_FIELDS;
+        }
+    }
+}