Fix bug that prevented some operations and service clients from resolving a service endpoint

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

@@ -15,17 +15,19 @@
 import software.amazon.smithy.go.codegen.integration.RuntimeClientPlugin;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.knowledge.ServiceIndex;
+import software.amazon.smithy.model.knowledge.TopDownIndex;
 import software.amazon.smithy.model.shapes.OperationShape;
 import software.amazon.smithy.model.shapes.ServiceShape;
 import software.amazon.smithy.model.shapes.ShapeId;
+import software.amazon.smithy.model.shapes.ToShapeId;
 import software.amazon.smithy.model.traits.Trait;
 import software.amazon.smithy.utils.ListUtils;
 
 public final class RegisterServiceMetadataMiddleware implements GoIntegration {
     List<RuntimeClientPlugin> runtimeClientPlugins = new ArrayList<>();
 
     @Override
-    public byte getOrder(){
+    public byte getOrder() {
         return 30;
     }
 
@@ -41,39 +43,49 @@ public void writeAdditionalFiles(
                 "RegisterServiceMetadata", AwsGoDependency.AWS_MIDDLEWARE).build();
         ServiceIndex serviceIndex = ServiceIndex.of(model);
 
-        for (ShapeId operationId: service.getAllOperations()) {
-            String middlewareName = getServiceMetadataMiddlewareName(service, operationId);
-            OperationShape operation = model.expectShape(operationId, OperationShape.class);
-            goDelegator.useShapeWriter(operation, writer -> {
+        TopDownIndex topDownIndex = TopDownIndex.of(model);
+
+        for (ToShapeId operation : topDownIndex.getContainedOperations(service)) {
+            String middlewareName = getServiceMetadataMiddlewareName(service, operation.toShapeId());
+            OperationShape operationShape = model.expectShape(operation.toShapeId(), OperationShape.class);
+            goDelegator.useShapeWriter(operationShape, writer -> {
                 writer.openBlock("func $L(region string) $P {", "}",
                         middlewareName, serviceMetadataProvider, () -> {
-                    StringBuilder builder = new StringBuilder();
-                    builder.append(" return &$T{\n");
-                    builder.append("Region: region,\n");
-                    builder.append("ServiceID: ServiceID,\n");
+                            StringBuilder builder = new StringBuilder();
+                            builder.append(" return &$T{\n");
+                            builder.append("Region: region,\n");
+                            builder.append("ServiceID: ServiceID,\n");
 
-                    Map<ShapeId, Trait> authSchemes = serviceIndex.getEffectiveAuthSchemes(service, operation);
-                    if (authSchemes.containsKey(SigV4Trait.ID)) {
-                        SigV4Trait trait = (SigV4Trait) authSchemes.get(SigV4Trait.ID);
-                        builder.append(String.format("SigningName: \"%s\",\n", trait.getName()));
-                    }
-                    builder.append(String.format("OperationName: \"%s\",\n", operation.getId().getName(service)));
-                    builder.append("}");
+                            Map<ShapeId, Trait> authSchemes = serviceIndex.getEffectiveAuthSchemes(service, operationShape);
+                            if (authSchemes.containsKey(SigV4Trait.ID)) {
+                                SigV4Trait trait = (SigV4Trait) authSchemes.get(SigV4Trait.ID);
+                                builder.append(String.format("SigningName: \"%s\",\n", trait.getName()));
+                            }
+                            builder.append(String.format("OperationName: \"%s\",\n",
+                                    operationShape.getId().getName(service)));
+                            builder.append("}");
 
-                    writer.write(builder.toString(), serviceMetadataProvider);
-                });
+                            writer.write(builder.toString(), serviceMetadataProvider);
+                        });
             });
         }
     }
 
-        @Override
+    @Override
     public void processFinalizedModel(GoSettings settings, Model model) {
         ServiceShape service = settings.getService(model);
-        for (ShapeId operationId : service.getAllOperations()) {
-            String middlewareName = getServiceMetadataMiddlewareName(service, operationId);
-            OperationShape operation = model.expectShape(operationId, OperationShape.class);
+        TopDownIndex index = TopDownIndex.of(model);
+
+        for (ToShapeId operation : index.getContainedOperations(service)) {
+            String middlewareName = getServiceMetadataMiddlewareName(service, operation.toShapeId());
+            OperationShape operationShape = model.expectShape(operation.toShapeId(), OperationShape.class);
             RuntimeClientPlugin runtimeClientPlugin = RuntimeClientPlugin.builder()
-                    .operationPredicate((m, s, o) -> operation.equals(o))
+                    .operationPredicate((m, s, o) -> {
+                        if (!s.equals(service)) {
+                            return false;
+                        }
+                        return operationShape.equals(o);
+                    })
                     .registerMiddleware(MiddlewareRegistrar.builder()
                             .resolvedFunction(SymbolUtils.createValueSymbolBuilder(
                                     middlewareName).build())
@@ -93,6 +105,6 @@ public List<RuntimeClientPlugin> getClientPlugins() {
     }
 
     private String getServiceMetadataMiddlewareName(ServiceShape service, ShapeId operationID) {
-        return "newServiceMetadataMiddleware_op"+ operationID.getName(service);
+        return "newServiceMetadataMiddleware_op" + operationID.getName(service);
     }
 }