Skip indexing endpoint modifier traits without shape definition

AndrewFossAWS · Jan 11, 2024 · 3609c8e · 3609c8e
1 parent 5b4bbae
commit 3609c8e
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 2 deletions.
diff --git a/...ts/src/main/java/software/amazon/smithy/rulesengine/aws/traits/EndpointModifierIndex.java b/...ts/src/main/java/software/amazon/smithy/rulesengine/aws/traits/EndpointModifierIndex.java
@@ -7,7 +7,10 @@
 
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Optional;
 import java.util.TreeMap;
+import java.util.logging.Level;
+import java.util.logging.Logger;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.knowledge.KnowledgeIndex;
 import software.amazon.smithy.model.shapes.ServiceShape;
@@ -22,15 +25,22 @@
  * Endpoint modifier traits are traits that are marked by {@link EndpointModifierTrait}
  */
 public final class EndpointModifierIndex implements KnowledgeIndex {
+    private static final Logger LOGGER = Logger.getLogger(EndpointModifierIndex.class.getName());
 
     private final Map<ShapeId, Map<ShapeId, Trait>> endpointModifierTraits = new HashMap<>();
 
     public EndpointModifierIndex(Model model) {
         for (ServiceShape serviceShape : model.getServiceShapes()) {
             Map<ShapeId, Trait> result = new TreeMap<>();
             for (Trait trait : serviceShape.getAllTraits().values()) {
-                Shape traitShape = model.expectShape(trait.toShapeId());
-                if (traitShape.hasTrait(EndpointModifierTrait.ID)) {
+                Optional<Shape> traitShape = model.getShape(trait.toShapeId());
+                if (!traitShape.isPresent()) {
+                    LOGGER.log(Level.WARNING, String.format(
+                        "%s trait found in service %s but it's missing definition",
+                            trait.toShapeId(), serviceShape.toShapeId()));
+                    continue;
+                }
+                if (traitShape.get().hasTrait(EndpointModifierTrait.ID)) {
                     result.put(trait.toShapeId(), trait);
                 }
             }

diff --git a/...rc/test/java/software/amazon/smithy/rulesengine/aws/traits/EndpointModifierIndexTest.java b/...rc/test/java/software/amazon/smithy/rulesengine/aws/traits/EndpointModifierIndexTest.java
@@ -46,6 +46,27 @@ public void loadsFromModel() {
             index.getEndpointModifierTraits(service4).get(RuleBasedEndpointsTrait.ID));
     }
 
+
+    @Test
+    public void indexSkipsLoadingTraitsWhenDefinitionIsMissing() {
+        Model model = Model.assembler()
+                .discoverModels(getClass().getClassLoader())
+                .addImport(getClass().getResource("endpointModifierIndex.smithy"))
+                .assemble()
+                .unwrap();
+
+        // Remove trait shape definition from model to verify index can skip it
+        model = model.toBuilder().removeShape(StandardRegionalEndpointsTrait.ID).build();
+
+        EndpointModifierIndex index = new EndpointModifierIndex(model);
+
+        ShapeId service1 = getServiceShapeId(model, "ns.foo#Service1");
+        ShapeId service3 = getServiceShapeId(model, "ns.foo#Service3");
+
+        assertEquals(index.getEndpointModifierTraits(service1).size(), 0);
+        assertEquals(index.getEndpointModifierTraits(service3).size(), 1);
+    }
+
     private ShapeId getServiceShapeId(Model model, String service) {
         return model
             .expectShape(ShapeId.from(service), ServiceShape.class).toShapeId();