Refactor JsonSchemaMapper to use JsonSchemaMapperContext

docs/source-2.0/guides/converting-to-openapi.rst

@@ -1841,4 +1841,4 @@ The conversion process is highly extensible through
 .. _x-amazon-apigateway-authorizer: https://docs.aws.amazon.com/apigateway/latest/developerguide/api-gateway-swagger-extensions-authorizer.html
 .. _Lambda authorizers: https://docs.aws.amazon.com/apigateway/latest/developerguide/api-gateway-swagger-extensions-authorizer.html
 .. _API Gateway's API key usage plans: https://docs.aws.amazon.com/apigateway/latest/developerguide/api-gateway-api-usage-plans.html
-.. _OpenAPI specification extension: https://spec.openapis.org/oas/v3.1.0#specification-extensions
+.. _OpenAPI specification extension: https://spec.openapis.org/oas/v3.1.0#specification-extensions

smithy-jsonschema/src/main/java/software/amazon/smithy/jsonschema/JsonSchemaMapper.java

@@ -19,6 +19,11 @@
 
 /**
  * Updates a schema builder before converting a shape to a schema.
+ *
+ * {@link JsonSchemaMapper#updateSchema(JsonSchemaMapperContext, Schema.Builder)} is the entry point during JSON Schema
+ * conversion, and is the recommended method to implement. If this method is implemented,
+ * {@link JsonSchemaMapper#updateSchema(Shape, Schema.Builder, JsonSchemaConfig)} will NOT be called unless written in
+ * the implementation.
  */
 public interface JsonSchemaMapper {
     /**
@@ -36,12 +41,33 @@ default byte getOrder() {
     }
 
     /**
-     * Updates a schema builder.
+     * Updates a schema builder using information in {@link JsonSchemaMapperContext}.
+     *
+     * If not implemented, will default to
+     * {@link JsonSchemaMapper#updateSchema(Shape, Schema.Builder, JsonSchemaConfig)} for backwards-compatibility.
+     *
+     * @param context Context with information needed to update the schema.
+     * @param schemaBuilder Schema builder to update.
+     * @return Returns an updated schema builder.
+     */
+    default Schema.Builder updateSchema(JsonSchemaMapperContext context, Schema.Builder schemaBuilder) {
+        return updateSchema(context.getShape(), schemaBuilder, context.getConfig());
+    }
+
+    /**
+     * Updates a schema builder, and is not recommended. Use
+     * {@link JsonSchemaMapper#updateSchema(JsonSchemaMapperContext, Schema.Builder)} instead.
+     *
+     * If not implemented, this method will default to a no-op.
+     *
+     * This method is not deprecated for backwards-compatibility.
      *
      * @param shape Shape used for the conversion.
      * @param schemaBuilder Schema builder to update.
      * @param config JSON Schema config.
      * @return Returns an updated schema builder.
      */
-    Schema.Builder updateSchema(Shape shape, Schema.Builder schemaBuilder, JsonSchemaConfig config);
+    default Schema.Builder updateSchema(Shape shape, Schema.Builder schemaBuilder, JsonSchemaConfig config) {
+        return schemaBuilder;
+    }
 }

smithy-jsonschema/src/main/java/software/amazon/smithy/jsonschema/JsonSchemaMapperContext.java

@@ -3,25 +3,22 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-
 package software.amazon.smithy.jsonschema;
 
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.shapes.Shape;
 
 /**
  * Context for a JSON schema mapping.
- *
- * @param <T> Type of Smithy {@link Shape} being mapped.
  */
-public class JsonSchemaMapperContext<T extends Shape> {
+public class JsonSchemaMapperContext {
     private final Model model;
-    private final T shape;
+    private final Shape shape;
     private final JsonSchemaConfig config;
 
     JsonSchemaMapperContext(
             Model model,
-            T shape,
+            Shape shape,
             JsonSchemaConfig config
     ) {
         this.model = model;
@@ -43,7 +40,7 @@ public Model getModel() {
      *
      * @return Returns the Smithy shape.
      */
-    public T getShape() {
+    public Shape getShape() {
         return shape;
     }
 

smithy-jsonschema/src/main/java/software/amazon/smithy/jsonschema/JsonSchemaShapeVisitor.java

@@ -343,16 +343,10 @@ private Schema.Builder updateBuilder(Shape shape, Schema.Builder builder) {
      * @param builder Schema being built.
      * @return Returns the built schema.
      */
-    private <T extends Shape> Schema buildSchema(T shape, Schema.Builder builder) {
-        JsonSchemaConfig config = converter.getConfig();
-        JsonSchemaMapperContext<T> context = new JsonSchemaMapperContext<>(model, shape, config);
-
+    private Schema buildSchema(Shape shape, Schema.Builder builder) {
+        JsonSchemaMapperContext context = new JsonSchemaMapperContext(model, shape, converter.getConfig());
         for (JsonSchemaMapper mapper : mappers) {
-            builder = mapper.updateSchema(shape, builder, config);
-
-            if (mapper instanceof JsonSchemaMapperV2) {
-                builder = ((JsonSchemaMapperV2) mapper).updateSchema(context, builder);
-            }
+            builder = mapper.updateSchema(context, builder);
         }
 
         return builder.build();

smithy-jsonschema/src/test/java/software/amazon/smithy/jsonschema/JsonSchemaConverterTest.java

@@ -211,7 +211,30 @@ public void canUseCustomPropertyNamingStrategy() {
     public void canAddCustomSchemaMapper() {
         Shape struct = StructureShape.builder().id("smithy.example#Foo").build();
         Model model = Model.builder().addShape(struct).build();
-        JsonSchemaMapper mapper = (shape, builder, conf) -> builder.putExtension("Hi", Node.from("There"));
+        class CustomMapper implements JsonSchemaMapper {
+            @Override
+            public Schema.Builder updateSchema(Shape shape, Schema.Builder builder, JsonSchemaConfig conf) {
+                return builder.putExtension("Hi", Node.from("There"));
+            }
+        }
+        JsonSchemaMapper mapper = new CustomMapper();
+        SchemaDocument doc = JsonSchemaConverter.builder().addMapper(mapper).model(model).build().convert();
+
+        assertTrue(doc.getDefinition("#/definitions/Foo").isPresent());
+        assertTrue(doc.getDefinition("#/definitions/Foo").get().getExtension("Hi").isPresent());
+    }
+
+    @Test
+    public void canAddCustomSchemaMapperContextMethod() {
+        Shape struct = StructureShape.builder().id("smithy.example#Foo").build();
+        Model model = Model.builder().addShape(struct).build();
+        class CustomMapper implements JsonSchemaMapper {
+            @Override
+            public Schema.Builder updateSchema(JsonSchemaMapperContext context, Schema.Builder builder) {
+                return builder.putExtension("Hi", Node.from("There"));
+            }
+        }
+        JsonSchemaMapper mapper = new CustomMapper();
         SchemaDocument doc = JsonSchemaConverter.builder().addMapper(mapper).model(model).build().convert();
 
         assertTrue(doc.getDefinition("#/definitions/Foo").isPresent());

smithy-openapi-traits/src/main/java/software/amazon/smithy/openapi/traits/SpecificationExtensionTrait.java

@@ -5,8 +5,9 @@
 
 package software.amazon.smithy.openapi.traits;
 
+import java.util.Optional;
 import software.amazon.smithy.model.node.Node;
-import software.amazon.smithy.model.node.ObjectNode;
+import software.amazon.smithy.model.node.NodeMapper;
 import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.traits.AbstractTrait;
 import software.amazon.smithy.model.traits.AbstractTraitBuilder;
@@ -17,14 +18,12 @@
  * <code>smithy.openapi#specificationExtension</code> - Indicates a trait shape should be converted into an <a href="https://spec.openapis.org/oas/v3.1.0#specification-extensions">OpenAPI specification extension</a>.
  */
 public final class SpecificationExtensionTrait extends AbstractTrait
-        implements ToSmithyBuilder<SpecificationExtensionTrait> {
+    implements ToSmithyBuilder<SpecificationExtensionTrait> {
     public static final ShapeId ID = ShapeId.from("smithy.openapi#specificationExtension");
 
-    private static final String AS_MEMBER_NAME = "as";
-
     private final String as;
 
-    private SpecificationExtensionTrait(SpecificationExtensionTrait.Builder builder) {
+    private SpecificationExtensionTrait(Builder builder) {
         super(ID, builder.getSourceLocation());
         this.as = builder.as;
     }
@@ -36,53 +35,47 @@ public Provider() {
 
         @Override
         public Trait createTrait(ShapeId target, Node value) {
-            ObjectNode node = value.expectObjectNode();
-            SpecificationExtensionTrait.Builder builder = builder().sourceLocation(value);
-            node.getStringMember(AS_MEMBER_NAME, builder::as);
-            SpecificationExtensionTrait trait = builder.build();
+            SpecificationExtensionTrait trait = new NodeMapper().deserialize(value, SpecificationExtensionTrait.class);
             trait.setNodeCache(value);
             return trait;
         }
     }
 
     /**
-     * Gets the extension name for a given trait shape.
-     * Either an explicitly configured extension name, or a default transformation of the shape ID.
+     * Gets the specification extension header value in "as".
      *
-     * @param traitShapeId Trait shape to get the extension name.
-     * @return Extension name for the given trait shape.
+     * @return Returns the optionally present "as".
      */
-    public String extensionNameFor(ShapeId traitShapeId) {
-        return as != null
-                ? as
-                : "x-" + traitShapeId.toString().replaceAll("[.#]", "-");
+    public Optional<String> getAs() {
+        return Optional.ofNullable(as);
     }
 
-    public static SpecificationExtensionTrait.Builder builder() {
-        return new SpecificationExtensionTrait.Builder();
+    public static Builder builder() {
+        return new Builder();
     }
 
     @Override
     protected Node createNode() {
-        return Node.objectNodeBuilder()
-                .sourceLocation(getSourceLocation())
-                .withMember(AS_MEMBER_NAME, this.as)
-                .build();
+        NodeMapper mapper = new NodeMapper();
+        mapper.disableToNodeForClass(SpecificationExtensionTrait.class);
+        mapper.setOmitEmptyValues(true);
+        return mapper.serialize(this).expectObjectNode();
     }
 
     @Override
-    public SpecificationExtensionTrait.Builder toBuilder() {
+    public Builder toBuilder() {
         return builder()
                 .sourceLocation(getSourceLocation())
                 .as(this.as);
     }
 
-    public static final class Builder
-            extends AbstractTraitBuilder<SpecificationExtensionTrait, SpecificationExtensionTrait.Builder> {
+    /**
+     * Builds a {@link SpecificationExtensionTrait} trait.
+     */
+    public static final class Builder extends AbstractTraitBuilder<SpecificationExtensionTrait, Builder> {
         private String as;
 
-        private Builder() {
-        }
+        private Builder() {}
 
         @Override
         public SpecificationExtensionTrait build() {
@@ -95,7 +88,7 @@ public SpecificationExtensionTrait build() {
          * @param as Explicit name for the target specification extension, or null.
          * @return This builder instance.
          */
-        public SpecificationExtensionTrait.Builder as(String as) {
+        public Builder as(String as) {
             this.as = as;
             return this;
         }

smithy-openapi-traits/src/main/resources/META-INF/smithy/manifest

@@ -1 +1 @@
-smithy.openapi.smithy
+smithy.openapi.smithy

smithy-openapi-traits/src/main/resources/META-INF/smithy/smithy.openapi.smithy

@@ -18,4 +18,4 @@ structure specificationExtension {
 
 @private
 @pattern("^x-.+$")
-string SpecificationExtensionKey
+string SpecificationExtensionKey

smithy-openapi/src/main/java/software/amazon/smithy/openapi/OpenApiUtils.java

@@ -0,0 +1,58 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.openapi;
+
+import java.util.LinkedHashMap;
+import java.util.Map;
+import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.node.Node;
+import software.amazon.smithy.model.shapes.Shape;
+import software.amazon.smithy.model.shapes.ShapeId;
+import software.amazon.smithy.openapi.traits.SpecificationExtensionTrait;
+import software.amazon.smithy.utils.SmithyInternalApi;
+
+@SmithyInternalApi
+public final class OpenApiUtils {
+    private OpenApiUtils() {}
+
+    /**
+     * Gets the specification extension name for a given meta trait.
+     *
+     * Either an explicitly configured extension name in specificationExtensionTrait, or a normalization of the shape
+     * ID. The normalization replaces all "." and "#" in a shapeId to "-".
+     *
+     * @param metaTraitId Trait shape to get the extension name.
+     * @return Extension name for the given trait shape.
+     */
+    public static String getSpecificationExtensionName(
+        ShapeId metaTraitId,
+        SpecificationExtensionTrait specificationExtensionTrait
+    ) {
+        return specificationExtensionTrait.getAs()
+            .orElse("x-" + metaTraitId.toString().replaceAll("[.#]", "-"));
+    }
+
+    /**
+     * Return specification extensions attached to a given shape.
+     *
+     * @param shape    Shape to get extensions for.
+     * @param model    Model the shape belongs to.
+     * @return map of specification extension names to node values
+     */
+    public static Map<String, Node> getSpecificationExtensionsMap(Model model, Shape shape) {
+        Map<String, Node> specificationExtensions = new LinkedHashMap<String, Node>();
+        shape.getAllTraits().forEach((traitId, trait) ->
+            // Get Applied Trait
+            model.getShape(traitId)
+                    // Get SpecificationExtensionTrait on the Applied Trait
+                    .flatMap(traitShape -> traitShape.getTrait(SpecificationExtensionTrait.class))
+                    // Get specification extension name from the Applied Trait and SpecificationExtensionTrait
+                    .map(specificationExtension -> getSpecificationExtensionName(traitId, specificationExtension))
+                    // Put the specification extension name and Applied Meta trait into the map.
+                    .ifPresent(name -> specificationExtensions.put(name, trait.toNode())));
+        return specificationExtensions;
+    }
+}

smithy-openapi/src/main/java/software/amazon/smithy/openapi/fromsmithy/OpenApiConverter.java

@@ -330,10 +330,7 @@ private <T extends Trait> OpenApi convertWithEnvironment(ConversionEnvironment<T
         openapi.components(environment.components.build());
 
         // Add arbitrary extensions if they're configured.
-        context.getConfig()
-                .getSchemaDocumentExtensions()
-                .getStringMap()
-                .forEach(openapi::putExtension);
+        openapi.getExtensions().putAll(context.getConfig().getSchemaDocumentExtensions().getStringMap());
 
         return mapper.after(context, openapi.build());
     }