fix(client-s3-control): codegen remove host prefix behavior

Author: kuhe

codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AddS3Config.java

@@ -19,7 +19,6 @@
 import static software.amazon.smithy.typescript.codegen.integration.RuntimeClientPlugin.Convention.HAS_CONFIG;
 import static software.amazon.smithy.typescript.codegen.integration.RuntimeClientPlugin.Convention.HAS_MIDDLEWARE;
 
-import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
@@ -44,9 +43,12 @@
 import software.amazon.smithy.model.shapes.TimestampShape;
 import software.amazon.smithy.model.traits.DeprecatedTrait;
 import software.amazon.smithy.model.traits.DocumentationTrait;
+import software.amazon.smithy.model.traits.EndpointTrait;
 import software.amazon.smithy.model.traits.HttpHeaderTrait;
 import software.amazon.smithy.model.traits.HttpPayloadTrait;
 import software.amazon.smithy.model.traits.StreamingTrait;
+import software.amazon.smithy.model.transform.ModelTransformer;
+import software.amazon.smithy.rulesengine.traits.EndpointRuleSetTrait;
 import software.amazon.smithy.typescript.codegen.LanguageTarget;
 import software.amazon.smithy.typescript.codegen.TypeScriptDependency;
 import software.amazon.smithy.typescript.codegen.TypeScriptSettings;
@@ -81,22 +83,16 @@ public final class AddS3Config implements TypeScriptIntegration {
         + "For more information, please go to https://github.com/aws/aws-sdk-js-v3#known-issues</p>";
 
     /**
-     * This removes the host prefix from the endpoint traits of the model.
-     * 
-    */
-    public static Model removeHostPrefixTrait(Model model, TypeScriptSettings settings) {
-       boolean checkEndPointRuleSet = model.hasTrait(EndpointRuleSetTrait.class)
-       Set<Shape> shapesWithTraits =  model.getShapesWithTrait(EndpointTrait.class);
-       Model.Builder mb = model.toBuilder();
-       for(Shape shape: shapesWithTraits) {
-            OperationShape.Builder ob = OperationShape.shapeToBuilder(shape.asOperationShape().get()); 
-            ob.removeTrait(EndpointTrait.ID);
-            OperationShape operation = ob.build();
-            mb.removeShape(operation.getId());
-            mb.addShape(operation);       
-       }
-       return mb.build();
-
+     * Remove the hostPrefix functionality by removing the endpoint traits.
+     * Only applied if endpointRuleSet trait present on S3/S3Control services.
+     */
+    public static Shape removeHostPrefixTrait(Shape shape) {
+        return shape.asOperationShape()
+            .map(OperationShape::shapeToBuilder)
+            .map(builder -> ((OperationShape.Builder) builder).removeTrait(EndpointTrait.ID))
+            .map(OperationShape.Builder::build)
+            .map(s -> (Shape) s)
+            .orElse(shape);
     }
 
     @Override
@@ -114,9 +110,9 @@ public Model preprocessModel(Model model, TypeScriptSettings settings) {
         if (!isS3(serviceShape)) {
             return model;
         }
-        removeHostPrefixTrait(model, settings);
 
         Model.Builder modelBuilder = model.toBuilder();
+        boolean hasRuleset = !model.getServiceShapesWithTrait(EndpointRuleSetTrait.class).isEmpty();
 
         TopDownIndex topDownIndex = TopDownIndex.of(model);
         Set<StructureShape> inputShapes = new HashSet<>();
@@ -227,7 +223,13 @@ public Model preprocessModel(Model model, TypeScriptSettings settings) {
             }
         }
 
-        return modelBuilder.addShapes(inputShapes).build();
+        Model builtModel = modelBuilder.addShapes(inputShapes).build();
+        if (hasRuleset) {
+            ModelTransformer.create().mapShapes(
+                builtModel, AddS3Config::removeHostPrefixTrait
+            );
+        }
+        return builtModel;
     }
 
     @Override

codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AddS3ControlDependency.java

@@ -23,6 +23,7 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.function.Consumer;
+import java.util.function.Function;
 import software.amazon.smithy.aws.traits.ServiceTrait;
 import software.amazon.smithy.codegen.core.SymbolProvider;
 import software.amazon.smithy.model.Model;
@@ -32,6 +33,7 @@
 import software.amazon.smithy.model.shapes.Shape;
 import software.amazon.smithy.model.traits.RequiredTrait;
 import software.amazon.smithy.model.transform.ModelTransformer;
+import software.amazon.smithy.rulesengine.traits.EndpointRuleSetTrait;
 import software.amazon.smithy.typescript.codegen.LanguageTarget;
 import software.amazon.smithy.typescript.codegen.TypeScriptSettings;
 import software.amazon.smithy.typescript.codegen.TypeScriptWriter;
@@ -90,16 +92,19 @@ public Model preprocessModel(Model model, TypeScriptSettings settings) {
         if (!isS3Control(settings.getService(model))) {
             return model;
         }
-        model = AddS3Config.removeHostPrefixTrait(model, settings);
+        boolean hasRuleset = !model.getServiceShapesWithTrait(EndpointRuleSetTrait.class).isEmpty();
         ServiceShape serviceShape = model.expectShape(settings.getService(), ServiceShape.class);
-        return ModelTransformer.create().mapShapes(model, shape -> {
+        Function<Shape, Shape> removeRequired = shape -> {
             Optional<MemberShape> modified = shape.asMemberShape()
-                    .filter(memberShape -> memberShape.getTarget().getName(serviceShape).equals("AccountId"))
-                    .filter(memberShape -> model.expectShape(memberShape.getTarget()).isStringShape())
-                    .filter(memberShape -> memberShape.isRequired())
-                    .map(memberShape -> Shape.shapeToBuilder(memberShape).removeTrait(RequiredTrait.ID).build());
+                .filter(memberShape -> memberShape.getTarget().getName(serviceShape).equals("AccountId"))
+                .filter(memberShape -> model.expectShape(memberShape.getTarget()).isStringShape())
+                .filter(MemberShape::isRequired)
+                .map(memberShape -> Shape.shapeToBuilder(memberShape).removeTrait(RequiredTrait.ID).build());
             return modified.isPresent() ? modified.get() : shape;
-        });
+        };
+        return ModelTransformer.create().mapShapes(
+            model, hasRuleset ? removeRequired.andThen(AddS3Config::removeHostPrefixTrait) : removeRequired
+        );
     }
 
     @Override