Repository: smithy-lang/smithy
Pull request #2096 (2024-01-11 19:02:32)

Skip indexing endpoint modifier traits without shape definition
*Description of changes:*
Skip indexing endpoint modifier traits without shape definition

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.



Patch:
diff --git a/smithy-aws-endpoints/src/main/java/software/amazon/smithy/rulesengine/aws/traits/EndpointModifierIndex.java b/smithy-aws-endpoints/src/main/java/software/amazon/smithy/rulesengine/aws/traits/EndpointModifierIndex.java
index ff8ccbc59..3d566876d 100644
--- a/smithy-aws-endpoints/src/main/java/software/amazon/smithy/rulesengine/aws/traits/EndpointModifierIndex.java
+++ b/smithy-aws-endpoints/src/main/java/software/amazon/smithy/rulesengine/aws/traits/EndpointModifierIndex.java
@@ -7,7 +7,9 @@ package software.amazon.smithy.rulesengine.aws.traits;
 
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Optional;
 import java.util.TreeMap;
+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,6 +24,7 @@ import software.amazon.smithy.model.traits.Trait;
  * 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<>();
 
@@ -29,8 +32,14 @@ public final class EndpointModifierIndex implements KnowledgeIndex {
         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.warning(String.format(
+                        "`%s` trait found in service `%s`, but the trait definition is missing",
+                            trait.toShapeId(), serviceShape.toShapeId()));
+                    continue;
+                }
+                if (traitShape.get().hasTrait(EndpointModifierTrait.ID)) {
                     result.put(trait.toShapeId(), trait);
                 }
             }


Test patch:
diff --git a/smithy-aws-endpoints/src/test/java/software/amazon/smithy/rulesengine/aws/traits/EndpointModifierIndexTest.java b/smithy-aws-endpoints/src/test/java/software/amazon/smithy/rulesengine/aws/traits/EndpointModifierIndexTest.java
index fa11edd67..f4da256f4 100644
--- a/smithy-aws-endpoints/src/test/java/software/amazon/smithy/rulesengine/aws/traits/EndpointModifierIndexTest.java
+++ b/smithy-aws-endpoints/src/test/java/software/amazon/smithy/rulesengine/aws/traits/EndpointModifierIndexTest.java
@@ -46,6 +46,27 @@ class EndpointModifierIndexTest {
             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();