Migrates integ test for s3, s3 control

aws/middleware/metadata.go

@@ -25,7 +25,7 @@ func (s RegisterServiceMetadata) HandleInitialize(
 	ctx context.Context, in middleware.InitializeInput, next middleware.InitializeHandler,
 ) (out middleware.InitializeOutput, metadata middleware.Metadata, err error) {
 	if len(s.ServiceID) > 0 {
-		ctx = setServiceID(ctx, s.ServiceID)
+		ctx = SetServiceID(ctx, s.ServiceID)
 	}
 	if len(s.SigningName) > 0 {
 		ctx = SetSigningName(ctx, s.SigningName)
@@ -46,6 +46,7 @@ type (
 	signingRegionKey struct{}
 	regionKey        struct{}
 	operationNameKey struct{}
+	partitionIDKey   struct{}
 )
 
 // GetServiceID retrieves the service id from the context.
@@ -78,6 +79,12 @@ func GetOperationName(ctx context.Context) (v string) {
 	return v
 }
 
+// GetPartitionID retrieves the endpoint partition id from the context.
+func GetPartitionID(ctx context.Context) string {
+	v, _ := ctx.Value(partitionIDKey{}).(string)
+	return v
+}
+
 // SetSigningName set or modifies the signing name on the context.
 func SetSigningName(ctx context.Context, value string) context.Context {
 	return context.WithValue(ctx, signingNameKey{}, value)
@@ -88,8 +95,8 @@ func SetSigningRegion(ctx context.Context, value string) context.Context {
 	return context.WithValue(ctx, signingRegionKey{}, value)
 }
 
-// setServiceID sets the service id on the context.
-func setServiceID(ctx context.Context, value string) context.Context {
+// SetServiceID sets the service id on the context.
+func SetServiceID(ctx context.Context, value string) context.Context {
 	return context.WithValue(ctx, serviceIDKey{}, value)
 }
 
@@ -102,3 +109,8 @@ func setRegion(ctx context.Context, value string) context.Context {
 func setOperationName(ctx context.Context, value string) context.Context {
 	return context.WithValue(ctx, operationNameKey{}, value)
 }
+
+// SetPartitionID sets the partition id of a resolved region on the context
+func SetPartitionID(ctx context.Context, value string) context.Context {
+	return context.WithValue(ctx, partitionIDKey{}, value)
+}

codegen/smithy-aws-go-codegen/src/main/java/software/amazon/smithy/aws/go/codegen/EndpointGenerator.java

@@ -43,7 +43,7 @@
  * Writes out a file that resolves endpoints using endpoints.json, but the
  * created resolver resolves endpoints for a single service.
  */
-final class EndpointGenerator implements Runnable {
+public class EndpointGenerator implements Runnable {
     public static final String MIDDLEWARE_NAME = "ResolveEndpoint";
     public static final String ADD_MIDDLEWARE_HELPER_NAME = String.format("add%sMiddleware", MIDDLEWARE_NAME);
     public static final String RESOLVER_INTERFACE_NAME = "EndpointResolver";
@@ -73,8 +73,11 @@ final class EndpointGenerator implements Runnable {
     private final ObjectNode endpointData;
     private final String endpointPrefix;
     private final Map<String, Partition> partitions = new TreeMap<>();
+    private final Boolean isInternalOnly;
+    private final String resolvedSdkID;
 
-    EndpointGenerator(
+
+    public EndpointGenerator(
             GoSettings settings,
             Model model,
             TriConsumer<String, String, Consumer<GoWriter>> writerFactory
@@ -87,6 +90,28 @@ final class EndpointGenerator implements Runnable {
         this.endpointData = Node.parse(IoUtils.readUtf8Resource(getClass(), "endpoints.json")).expectObjectNode();
         validateVersion();
         loadPartitions();
+        this.isInternalOnly = false;
+        this.resolvedSdkID = serviceShape.expectTrait(ServiceTrait.class).getSdkId();
+    }
+
+    public EndpointGenerator(
+            GoSettings settings,
+            Model model,
+            TriConsumer<String, String, Consumer<GoWriter>> writerFactory,
+            String sdkID,
+            String arnNamespace,
+            Boolean internalOnly
+    ) {
+        this.settings = settings;
+        this.model = model;
+        this.writerFactory = writerFactory;
+        serviceShape = settings.getService(model);
+        this.endpointPrefix = getEndpointPrefix(sdkID, arnNamespace);
+        this.endpointData = Node.parse(IoUtils.readUtf8Resource(getClass(), "endpoints.json")).expectObjectNode();
+        validateVersion();
+        loadPartitions();
+        this.isInternalOnly = internalOnly;
+        this.resolvedSdkID = sdkID;
     }
 
     private void validateVersion() {
@@ -105,6 +130,12 @@ private String getEndpointPrefix(ServiceShape service) {
         return endpointPrefixData.getStringMemberOrDefault(serviceTrait.getSdkId(), serviceTrait.getArnNamespace());
     }
 
+    private String getEndpointPrefix(String sdkId, String arnNamespace){
+        ObjectNode endpointPrefixData = Node.parse(IoUtils.readUtf8Resource(getClass(), "endpoint-prefix.json"))
+                .expectObjectNode();
+        return endpointPrefixData.getStringMemberOrDefault(sdkId, arnNamespace);
+    }
+
     private void loadPartitions() {
         List<ObjectNode> partitionObjects = endpointData
                 .expectArrayMember("partitions")
@@ -118,22 +149,30 @@ private void loadPartitions() {
 
     @Override
     public void run() {
-        writerFactory.accept("endpoints.go", settings.getModuleName(), writer -> {
-            generatePublicResolverTypes(writer);
-            generateMiddleware(writer);
-            generateAwsEndpointResolverWrapper(writer);
-        });
-        writerFactory.accept(INTERNAL_ENDPOINT_PACKAGE + "/endpoints.go", getInternalEndpointImportPath(), (writer) -> {
+        if (!this.isInternalOnly) {
+            writerFactory.accept("endpoints.go", settings.getModuleName(), writer -> {
+                generatePublicResolverTypes(writer);
+                generateMiddleware(writer);
+                generateAwsEndpointResolverWrapper(writer);
+            });
+        }
+
+        String pkgName = isInternalOnly? INTERNAL_ENDPOINT_PACKAGE + "/"+this.endpointPrefix :INTERNAL_ENDPOINT_PACKAGE;
+        writerFactory.accept(pkgName + "/endpoints.go", getInternalEndpointImportPath(), (writer) -> {
             generateInternalResolverImplementation(writer);
             generateInternalEndpointsModel(writer);
         });
-        writerFactory.accept(INTERNAL_ENDPOINT_PACKAGE + "/endpoints_test.go",
-                getInternalEndpointImportPath(), (writer) -> {
-                    writer.addUseImports(SmithyGoDependency.TESTING);
-                    writer.openBlock("func TestRegexCompile(t *testing.T) {", "}", () -> {
-                        writer.write("_ = $T", getInternalEndpointsSymbol(INTERNAL_ENDPOINTS_DATA_NAME, false).build());
-                    });
-                });
+
+        if (!this.isInternalOnly) {
+            writerFactory.accept(INTERNAL_ENDPOINT_PACKAGE + "/endpoints_test.go",
+                    getInternalEndpointImportPath(), (writer) -> {
+                        writer.addUseImports(SmithyGoDependency.TESTING);
+                        writer.openBlock("func TestRegexCompile(t *testing.T) {", "}", () -> {
+                            writer.write("_ = $T", getInternalEndpointsSymbol(INTERNAL_ENDPOINTS_DATA_NAME, false).build());
+                        });
+            });
+        }
+
     }
 
     private void generateAwsEndpointResolverWrapper(GoWriter writer) {
@@ -252,12 +291,13 @@ private void generateMiddlewareResolverBody(GoStackStepMiddlewareGenerator g, Go
             });
             w.write("ctx = awsmiddleware.SetSigningName(ctx, signingName)");
         });
-        w.write("");
-
-        w.write("ctx = awsmiddleware.SetSigningRegion(ctx, endpoint.SigningRegion)");
         w.write("ctx = smithyhttp.SetHostnameImmutable(ctx, endpoint.HostnameImmutable)");
-        w.write("");
+        // set signing region on context
+        w.write("ctx = awsmiddleware.SetSigningRegion(ctx, endpoint.SigningRegion)");
+        // set partition id on context
+        w.write("ctx = awsmiddleware.SetPartitionID(ctx, endpoint.PartitionID)");
 
+        w.insertTrailingNewline();
         w.write("return next.HandleSerialize(ctx, in)");
     }
 
@@ -387,8 +427,10 @@ private void generateInternalResolverImplementation(GoWriter writer) {
 
         // Resolver
         Symbol resolverImplSymbol = SymbolUtils.createPointableSymbolBuilder(INTERNAL_RESOLVER_NAME).build();
+
+
         writer.writeDocs(String.format("%s %s endpoint resolver", resolverImplSymbol.getName(),
-                serviceShape.expectTrait(ServiceTrait.class).getSdkId()));
+                this.resolvedSdkID));
         writer.openBlock("type $T struct {", "}", resolverImplSymbol, () -> {
             writer.write("partitions $T", SymbolUtils.createValueSymbolBuilder("Partitions",
                     AwsGoDependency.AWS_ENDPOINTS).build());

codegen/smithy-aws-go-codegen/src/main/java/software/amazon/smithy/aws/go/codegen/XMLProtocolUtils.java

@@ -10,6 +10,7 @@
 import software.amazon.smithy.codegen.core.SymbolProvider;
 import software.amazon.smithy.go.codegen.GoWriter;
 import software.amazon.smithy.go.codegen.SmithyGoDependency;
+import software.amazon.smithy.go.codegen.SymbolUtils;
 import software.amazon.smithy.go.codegen.SyntheticClone;
 import software.amazon.smithy.go.codegen.integration.ProtocolGenerator;
 import software.amazon.smithy.model.Model;
@@ -274,12 +275,28 @@ public static void writeXmlErrorMessageCodeDeserializer(ProtocolGenerator.Genera
         ServiceShape service = context.getService();
 
         if (requiresS3Customization(service)) {
-            writer.addUseImports(AwsCustomGoDependency.S3_SHARED_CUSTOMIZATION);
+            Symbol getErrorComponentFunction = SymbolUtils.createValueSymbolBuilder(
+                    "GetErrorResponseComponents",
+                    AwsCustomGoDependency.S3_SHARED_CUSTOMIZATION
+            ).build();
+
+            Symbol errorOptions = SymbolUtils.createValueSymbolBuilder(
+                    "ErrorResponseDeserializerOptions",
+                    AwsCustomGoDependency.S3_SHARED_CUSTOMIZATION
+            ).build();
+
             if (isS3Service(service)){
-                writer.write("errorComponents, err := s3shared.GetS3ErrorResponseComponents(errorBody, response.StatusCode)");
+                // s3 service
+                writer.openBlock("errorComponents, err := $T(errorBody, $T{",
+                        "})", getErrorComponentFunction, errorOptions, () -> {
+                    writer.write("UseStatusCode : true, StatusCode : response.StatusCode,");
+                });
             } else {
                 // s3 control
-                writer.write("errorComponents, err := s3shared.GetErrorResponseComponents(errorBody)");
+                writer.openBlock("errorComponents, err := $T(errorBody, $T{",
+                        "})", getErrorComponentFunction, errorOptions, () -> {
+                            writer.write("IsWrappedWithErrorTag: true,");
+                        });
             }
 
             writer.write("if err != nil { return err }");

codegen/smithy-aws-go-codegen/src/main/java/software/amazon/smithy/aws/go/codegen/customization/S3ControlEndpointResolver.java

@@ -0,0 +1,34 @@
+package software.amazon.smithy.aws.go.codegen.customization;
+
+import java.util.function.Consumer;
+import software.amazon.smithy.aws.go.codegen.EndpointGenerator;
+import software.amazon.smithy.aws.traits.ServiceTrait;
+import software.amazon.smithy.codegen.core.SymbolProvider;
+import software.amazon.smithy.go.codegen.GoSettings;
+import software.amazon.smithy.go.codegen.GoWriter;
+import software.amazon.smithy.go.codegen.TriConsumer;
+import software.amazon.smithy.go.codegen.integration.GoIntegration;
+import software.amazon.smithy.model.Model;
+
+/**
+ * S3ControlEndpointResolverCustomizations adds an internal endpoint resolver
+ * for s3 service endpoints
+ */
+public class S3ControlEndpointResolver implements GoIntegration  {
+
+    @Override
+    public void writeAdditionalFiles(
+            GoSettings settings,
+            Model model,
+            SymbolProvider symbolProvider,
+            TriConsumer<String, String, Consumer<GoWriter>> writerFactory
+    ) {
+        if (!settings.getService(model).expectTrait(ServiceTrait.class).getSdkId().equalsIgnoreCase(
+                "S3 Control")){
+           return;
+        }
+
+        // Generate S3 internal endpoint resolver for S3 Control service
+        new EndpointGenerator(settings, model, writerFactory,"S3","s3", true).run();
+    }
+}