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 ff8ccbc59f1..006ec7f3312 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,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,6 +25,7 @@ * 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> endpointModifierTraits = new HashMap<>(); @@ -29,8 +33,14 @@ public EndpointModifierIndex(Model model) { for (ServiceShape serviceShape : model.getServiceShapes()) { Map result = new TreeMap<>(); for (Trait trait : serviceShape.getAllTraits().values()) { - Shape traitShape = model.expectShape(trait.toShapeId()); - if (traitShape.hasTrait(EndpointModifierTrait.ID)) { + Optional 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/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 fa11edd67aa..f4da256f486 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 @@ 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();