diff --git a/smithy-aws-traits/src/main/java/software/amazon/smithy/aws/traits/clientendpointdiscovery/CleanClientDiscoveryTraitTransformer.java b/smithy-aws-traits/src/main/java/software/amazon/smithy/aws/traits/clientendpointdiscovery/CleanClientDiscoveryTraitTransformer.java index 313888ddec8..16a3ca50362 100644 --- a/smithy-aws-traits/src/main/java/software/amazon/smithy/aws/traits/clientendpointdiscovery/CleanClientDiscoveryTraitTransformer.java +++ b/smithy-aws-traits/src/main/java/software/amazon/smithy/aws/traits/clientendpointdiscovery/CleanClientDiscoveryTraitTransformer.java @@ -15,7 +15,6 @@ import software.amazon.smithy.model.shapes.ServiceShape; import software.amazon.smithy.model.shapes.Shape; import software.amazon.smithy.model.shapes.ShapeId; -import software.amazon.smithy.model.traits.ErrorTrait; import software.amazon.smithy.model.transform.ModelTransformer; import software.amazon.smithy.model.transform.ModelTransformerPlugin; import software.amazon.smithy.utils.SmithyInternalApi; @@ -27,16 +26,18 @@ public final class CleanClientDiscoveryTraitTransformer implements ModelTransformerPlugin { @Override public Model onRemove(ModelTransformer transformer, Collection shapes, Model model) { + // Short circuit if the traits aren't applied. + if (!model.isTraitApplied(ClientEndpointDiscoveryTrait.class) + && !model.isTraitApplied(ClientDiscoveredEndpointTrait.class)) { + return model; + } + Set removedOperations = shapes.stream() .filter(Shape::isOperationShape) .map(Shape::getId) .collect(Collectors.toSet()); - Set removedErrors = shapes.stream() - .filter(shape -> shape.hasTrait(ErrorTrait.class)) - .map(Shape::getId) - .collect(Collectors.toSet()); - Set servicesToUpdate = getServicesToUpdate(model, removedOperations, removedErrors); + Set servicesToUpdate = getServicesToUpdate(model, removedOperations); Set shapesToUpdate = new HashSet<>(servicesToUpdate); Set operationsToUpdate = getOperationsToUpdate( @@ -51,7 +52,7 @@ public Model onRemove(ModelTransformer transformer, Collection shapes, Mo return transformer.replaceShapes(model, shapesToUpdate); } - private Set getServicesToUpdate(Model model, Set removedOperations, Set removedErrors) { + private Set getServicesToUpdate(Model model, Set removedOperations) { Set result = new HashSet<>(); for (ServiceShape service : model.getServiceShapesWithTrait(ClientEndpointDiscoveryTrait.class)) { ClientEndpointDiscoveryTrait trait = service.expectTrait(ClientEndpointDiscoveryTrait.class);