Add endpoint discovery traits

aws/smithy-aws-traits/src/main/java/software/amazon/smithy/aws/traits/endpointdiscovery/CleanDiscoveryTraitTransformer.java

@@ -0,0 +1,61 @@
+/*
+ * Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.smithy.aws.traits.endpointdiscovery;
+
+import java.util.Collection;
+import java.util.Set;
+import java.util.stream.Collectors;
+import software.amazon.smithy.model.Model;
+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.traits.Trait;
+import software.amazon.smithy.model.transform.ModelTransformer;
+import software.amazon.smithy.model.transform.ModelTransformerPlugin;
+
+/**
+ * Removes the endpoint discovery trait from a service if the referenced operation or error are removed.
+ */
+public class CleanDiscoveryTraitTransformer implements ModelTransformerPlugin {
+    @Override
+    public Model onRemove(ModelTransformer transformer, Collection<Shape> shapes, Model model) {
+        Set<ShapeId> removedOperations = shapes.stream()
+                .filter(Shape::isOperationShape)
+                .map(Shape::getId)
+                .collect(Collectors.toSet());
+        Set<ShapeId> removedErrors = shapes.stream()
+                .filter(shape -> shape.hasTrait(ErrorTrait.class))
+                .map(Shape::getId)
+                .collect(Collectors.toSet());
+
+        Set<Shape> toReplace = getServicesToUpdate(model, removedOperations, removedErrors);
+        return transformer.replaceShapes(model, toReplace);
+    }
+
+    private Set<Shape> getServicesToUpdate(Model model, Set<ShapeId> removedOperations, Set<ShapeId> removedErrors) {
+        return model.getShapeIndex().shapes(ServiceShape.class)
+                .flatMap(service -> Trait.flatMapStream(service, EndpointDiscoveryTrait.class))
+                .filter(pair -> removedOperations.contains(pair.getRight().getOperation())
+                        || removedErrors.contains(pair.getRight().getError()))
+                .map(pair -> {
+                    ServiceShape.Builder builder = pair.getLeft().toBuilder();
+                    builder.removeTrait(EndpointDiscoveryTrait.ID);
+                    return builder.build();
+                })
+                .collect(Collectors.toSet());
+    }
+}

aws/smithy-aws-traits/src/main/java/software/amazon/smithy/aws/traits/endpointdiscovery/DiscoveredEndpointTrait.java

@@ -0,0 +1,99 @@
+/*
+ * Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.smithy.aws.traits.endpointdiscovery;
+
+import java.util.Collections;
+import software.amazon.smithy.model.node.Node;
+import software.amazon.smithy.model.node.ObjectNode;
+import software.amazon.smithy.model.shapes.ShapeId;
+import software.amazon.smithy.model.traits.AbstractTrait;
+import software.amazon.smithy.model.traits.AbstractTraitBuilder;
+import software.amazon.smithy.utils.SmithyBuilder;
+import software.amazon.smithy.utils.ToSmithyBuilder;
+
+/**
+ * Indicates that the target operation should use the SDK's endpoint discovery
+ * logic.
+ */
+public final class DiscoveredEndpointTrait extends AbstractTrait implements ToSmithyBuilder<DiscoveredEndpointTrait> {
+    public static final ShapeId ID = ShapeId.from("aws.api#discoveredEndpoint");
+
+    private static final String REQUIRED = "required";
+
+    private final boolean required;
+
+    private DiscoveredEndpointTrait(Builder builder) {
+        super(ID, builder.getSourceLocation());
+        this.required = builder.required;
+    }
+
+    /**
+     * @return Returns a builder used to create {@link DiscoveredEndpointTrait}
+     */
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    /**
+     * @return Returns whether or not the service requires endpoint discovery.
+     */
+    public boolean isRequired() {
+        return required;
+    }
+
+    @Override
+    protected Node createNode() {
+        return Node.objectNode().withMember(REQUIRED, Node.from(isRequired()));
+    }
+
+    @Override
+    public SmithyBuilder<DiscoveredEndpointTrait> toBuilder() {
+        return builder().required(required);
+    }
+
+    /** Builder for {@link DiscoveredEndpointTrait}. */
+    public static final class Builder extends AbstractTraitBuilder<DiscoveredEndpointTrait, Builder> {
+        private boolean required;
+
+        private Builder() {}
+
+        @Override
+        public DiscoveredEndpointTrait build() {
+            return new DiscoveredEndpointTrait(this);
+        }
+
+        public Builder required(boolean required) {
+            this.required = required;
+            return this;
+        }
+    }
+
+    public static final class Provider extends AbstractTrait.Provider {
+        public Provider() {
+            super(ID);
+        }
+
+        @Override
+        public DiscoveredEndpointTrait createTrait(ShapeId target, Node value) {
+            ObjectNode objectNode = value.expectObjectNode();
+            objectNode.warnIfAdditionalProperties(Collections.singletonList(REQUIRED));
+
+            return builder()
+                    .required(objectNode.getBooleanMemberOrDefault(REQUIRED, true))
+                    .build();
+        }
+    }
+}

aws/smithy-aws-traits/src/main/java/software/amazon/smithy/aws/traits/endpointdiscovery/EndpointDiscoveryIdTrait.java

@@ -0,0 +1,42 @@
+/*
+ * Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.smithy.aws.traits.endpointdiscovery;
+
+import software.amazon.smithy.model.SourceLocation;
+import software.amazon.smithy.model.shapes.ShapeId;
+import software.amazon.smithy.model.traits.BooleanTrait;
+
+/**
+ * Indicates members of the operation input which should be use to discover
+ * endpoints.
+ */
+public final class EndpointDiscoveryIdTrait extends BooleanTrait {
+    public static final ShapeId ID = ShapeId.from("aws.api#endpointDiscoveryId");
+
+    public EndpointDiscoveryIdTrait(SourceLocation sourceLocation) {
+        super(ID, sourceLocation);
+    }
+
+    public EndpointDiscoveryIdTrait() {
+        this(SourceLocation.NONE);
+    }
+
+    public static final class Provider extends BooleanTrait.Provider<EndpointDiscoveryIdTrait> {
+        public Provider() {
+            super(ID, EndpointDiscoveryIdTrait::new);
+        }
+    }
+}

aws/smithy-aws-traits/src/main/java/software/amazon/smithy/aws/traits/endpointdiscovery/EndpointDiscoveryIndex.java

@@ -0,0 +1,114 @@
+/*
+ * Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.smithy.aws.traits.endpointdiscovery;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.stream.Collectors;
+import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.knowledge.KnowledgeIndex;
+import software.amazon.smithy.model.knowledge.OperationIndex;
+import software.amazon.smithy.model.knowledge.TopDownIndex;
+import software.amazon.smithy.model.shapes.MemberShape;
+import software.amazon.smithy.model.shapes.OperationShape;
+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.shapes.ShapeIndex;
+import software.amazon.smithy.model.shapes.StructureShape;
+import software.amazon.smithy.model.shapes.ToShapeId;
+import software.amazon.smithy.model.traits.Trait;
+import software.amazon.smithy.utils.Pair;
+
+public final class EndpointDiscoveryIndex implements KnowledgeIndex {
+    private final Map<ShapeId, Map<ShapeId, EndpointDiscoveryInfo>> endpointDiscoveryInfo = new HashMap<>();
+
+    public EndpointDiscoveryIndex(Model model) {
+        ShapeIndex index = model.getShapeIndex();
+        TopDownIndex topDownIndex = model.getKnowledge(TopDownIndex.class);
+        OperationIndex opIndex = model.getKnowledge(OperationIndex.class);
+
+        index.shapes(ServiceShape.class)
+                .flatMap(service -> Trait.flatMapStream(service, EndpointDiscoveryTrait.class))
+                .forEach(servicePair -> {
+                    ServiceShape service = servicePair.getLeft();
+                    ShapeId endpointOperationId = servicePair.getRight().getOperation();
+                    ShapeId endpointErrorId = servicePair.getRight().getError();
+
+                    Optional<OperationShape> endpointOperation = index.getShape(endpointOperationId)
+                            .flatMap(Shape::asOperationShape);
+                    Optional<StructureShape> endpointError = index.getShape(endpointErrorId)
+                            .flatMap(Shape::asStructureShape);
+
+                    if (endpointOperation.isPresent() && endpointError.isPresent()) {
+                        Map<ShapeId, EndpointDiscoveryInfo> serviceInfo = getOperations(
+                                service, endpointOperation.get(), endpointError.get(), topDownIndex, opIndex);
+                        if (!serviceInfo.isEmpty()) {
+                            endpointDiscoveryInfo.put(service.getId(), serviceInfo);
+                        }
+                    }
+                });
+    }
+
+    private Map<ShapeId, EndpointDiscoveryInfo> getOperations(
+            ServiceShape service,
+            OperationShape endpointOperation,
+            StructureShape endpointError,
+            TopDownIndex topDownIndex,
+            OperationIndex opIndex
+    ) {
+        return topDownIndex.getContainedOperations(service).stream()
+                .flatMap(operation -> Trait.flatMapStream(operation, DiscoveredEndpointTrait.class))
+                .map(pair -> {
+                    OperationShape operation = pair.getLeft();
+                    List<MemberShape> discoveryIds = getDiscoveryIds(opIndex, operation);
+                    EndpointDiscoveryInfo info = new EndpointDiscoveryInfo(
+                            service,
+                            operation,
+                            endpointOperation,
+                            endpointError,
+                            discoveryIds,
+                            pair.getRight().isRequired()
+                    );
+                    return Pair.of(operation.getId(), info);
+                })
+                .collect(Collectors.toMap(Pair::getLeft, Pair::getRight));
+    }
+
+    private List<MemberShape> getDiscoveryIds(OperationIndex opIndex, OperationShape operation) {
+        List<MemberShape> discoveryIds = new ArrayList<>();
+        opIndex.getInput(operation).ifPresent(input -> input.getAllMembers().values().stream()
+                .filter(member -> member.hasTrait(EndpointDiscoveryIdTrait.class))
+                .forEach(discoveryIds::add));
+        return discoveryIds;
+    }
+
+    public Optional<EndpointDiscoveryInfo> getEndpointDiscoveryInfo(ToShapeId service, ToShapeId operation) {
+        return Optional.ofNullable(endpointDiscoveryInfo.get(service.toShapeId()))
+                .flatMap(mappings -> Optional.ofNullable(mappings.get(operation.toShapeId())));
+    }
+
+    public Set<ShapeId> getEndpointDiscoveryOperations(ToShapeId service) {
+        return Optional.ofNullable(endpointDiscoveryInfo.get(service.toShapeId()))
+                .flatMap(mappings -> Optional.of(mappings.keySet()))
+                .orElse(new HashSet<>());
+    }
+}

aws/smithy-aws-traits/src/main/java/software/amazon/smithy/aws/traits/endpointdiscovery/EndpointDiscoveryInfo.java

@@ -0,0 +1,72 @@
+/*
+ * Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.smithy.aws.traits.endpointdiscovery;
+
+import java.util.List;
+import software.amazon.smithy.model.shapes.MemberShape;
+import software.amazon.smithy.model.shapes.OperationShape;
+import software.amazon.smithy.model.shapes.ServiceShape;
+import software.amazon.smithy.model.shapes.StructureShape;
+
+public final class EndpointDiscoveryInfo {
+
+    private final ServiceShape service;
+    private final OperationShape operation;
+    private final OperationShape discoveryOperation;
+    private final StructureShape error;
+    private final List<MemberShape> discoveryIds;
+    private final boolean required;
+
+    EndpointDiscoveryInfo(
+            ServiceShape service,
+            OperationShape operation,
+            OperationShape discoveryOperation,
+            StructureShape error,
+            List<MemberShape> discoveryIds,
+            boolean required
+    ) {
+        this.service = service;
+        this.operation = operation;
+        this.discoveryOperation = discoveryOperation;
+        this.error = error;
+        this.discoveryIds = discoveryIds;
+        this.required = required;
+    }
+
+    public ServiceShape getService() {
+        return service;
+    }
+
+    public OperationShape getOperation() {
+        return operation;
+    }
+
+    public OperationShape getDiscoveryOperation() {
+        return discoveryOperation;
+    }
+
+    public StructureShape getError() {
+        return error;
+    }
+
+    public List<MemberShape> getDiscoveryIds() {
+        return discoveryIds;
+    }
+
+    public boolean isRequired() {
+        return required;
+    }
+}