Cleanup how dynamic schemas work

Rather than make WrappedDocument and SchemaGuidedDocumentBuilder
support all kinds of shape types, they now just support structure and
document. This cleans up the handling of document and nested members
so that nested structures don't have to be re-parsed each time we need
to serialize them. It's also easier to figure out how the code works.

Various bugs were fixed along the way and more test coverage was
added.

Author: mtdowling

aws/sdkv2/aws-sdkv2-shapes/src/test/java/software/amazon/smithy/java/aws/sdkv2/shapes/SdkDocumentWriterTest.java

@@ -124,7 +124,7 @@ public static List<Arguments> restJsonTestCases() {
                                         "bar",
                                         software.amazon.awssdk.core.document.Document.fromString("a"),
                                         "BAZ",
-                                        software.amazon.awssdk.core.document.Document.fromString("b"),
+                                        software.amazon.awssdk.core.document.Document.fromNumber(2),
                                         "date",
                                         software.amazon.awssdk.core.document.Document
                                                 .fromString(Instant.EPOCH.toString())))));
@@ -150,7 +150,7 @@ private static SerializableStruct createFoo(DynamicClient client) {
                                 "bar",
                                 Document.of("a"),
                                 "baz",
-                                Document.of("b"),
+                                Document.of(2),
                                 "date",
                                 Document.of(Instant.EPOCH))));
     }
@@ -170,7 +170,7 @@ public void awsJsonDoesNotUseJsonName() {
                                         "bar",
                                         software.amazon.awssdk.core.document.Document.fromString("a"),
                                         "baz",
-                                        software.amazon.awssdk.core.document.Document.fromString("b"),
+                                        software.amazon.awssdk.core.document.Document.fromNumber(2),
                                         "date",
                                         software.amazon.awssdk.core.document.Document.fromNumber(0.0)))));
     }

client/client-core/src/test/java/software/amazon/smithy/java/client/core/plugins/InjectIdempotencyTokenPluginTest.java

@@ -99,7 +99,7 @@ private String callAndGetToken(String operation, Document input) {
                 .addPlugin(mock)
                 .build();
 
-        client.call(operation, Document.of(input));
+        client.call(operation, input);
         assertThat(mock.getRequests(), not(empty()));
         return mock.getRequests().get(0).request().headers().firstValue("x-token");
     }

client/dynamic-client/src/main/java/software/amazon/smithy/java/dynamicclient/DocumentException.java

@@ -12,17 +12,17 @@
 import software.amazon.smithy.java.core.serde.ShapeSerializer;
 import software.amazon.smithy.java.core.serde.document.Document;
 import software.amazon.smithy.java.dynamicschemas.SchemaConverter;
-import software.amazon.smithy.java.dynamicschemas.WrappedDocument;
+import software.amazon.smithy.java.dynamicschemas.StructDocument;
 import software.amazon.smithy.model.shapes.ShapeId;
 
 /**
  * A {@link ModeledException} that provides access to the contents of the exception as a document.
  */
 public final class DocumentException extends ModeledException {
 
-    private final WrappedDocument document;
+    private final StructDocument document;
 
-    DocumentException(Schema schema, String message, WrappedDocument document) {
+    DocumentException(Schema schema, String message, StructDocument document) {
         super(schema, message);
         this.document = document;
     }
@@ -54,7 +54,7 @@ public Document getContents() {
     public static final class SchemaGuidedExceptionBuilder implements ShapeBuilder<ModeledException> {
 
         private final Schema target;
-        private final ShapeBuilder<WrappedDocument> delegateBuilder;
+        private final ShapeBuilder<StructDocument> delegateBuilder;
 
         public SchemaGuidedExceptionBuilder(ShapeId service, Schema target) {
             this.target = target;

client/dynamic-client/src/main/java/software/amazon/smithy/java/dynamicclient/DynamicClient.java

@@ -31,14 +31,13 @@
 import software.amazon.smithy.java.core.serde.TypeRegistry;
 import software.amazon.smithy.java.core.serde.document.Document;
 import software.amazon.smithy.java.dynamicschemas.SchemaConverter;
-import software.amazon.smithy.java.dynamicschemas.WrappedDocument;
+import software.amazon.smithy.java.dynamicschemas.StructDocument;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.knowledge.ServiceIndex;
 import software.amazon.smithy.model.knowledge.TopDownIndex;
 import software.amazon.smithy.model.shapes.OperationShape;
 import software.amazon.smithy.model.shapes.ServiceShape;
 import software.amazon.smithy.model.shapes.ShapeId;
-import software.amazon.smithy.model.shapes.ShapeType;
 import software.amazon.smithy.model.shapes.ToShapeId;
 
 /**
@@ -65,7 +64,7 @@ public final class DynamicClient extends Client {
 
     private final ServiceShape service;
     private final Model model;
-    private final ConcurrentMap<String, ApiOperation<WrappedDocument, WrappedDocument>> operations =
+    private final ConcurrentMap<String, ApiOperation<StructDocument, StructDocument>> operations =
             new ConcurrentHashMap<>();
     private final SchemaConverter schemaConverter;
     private final Map<String, OperationShape> operationNames = new HashMap<>();
@@ -179,7 +178,7 @@ public CompletableFuture<Document> callAsync(
             RequestOverrideConfig overrideConfig
     ) {
         var apiOperation = getApiOperation(operation);
-        var inputStruct = new WrappedDocument(apiOperation.inputSchema(), input, service.getId());
+        var inputStruct = StructDocument.of(apiOperation.inputSchema(), input, service.getId());
         return call(inputStruct, apiOperation, overrideConfig).thenApply(Function.identity());
     }
 
@@ -192,13 +191,10 @@ public CompletableFuture<Document> callAsync(
      */
     public SerializableStruct createStruct(ToShapeId shape, Document value) {
         var schema = schemaConverter.getSchema(model.expectShape(shape.toShapeId()));
-        if (value.type() != ShapeType.MAP && value.type() != ShapeType.STRUCTURE) {
-            throw new IllegalArgumentException("Document value must be a map or structure, found " + value.type());
-        }
-        return new WrappedDocument(schema, value, service.getId());
+        return StructDocument.of(schema, value, service.getId());
     }
 
-    private ApiOperation<WrappedDocument, WrappedDocument> getApiOperation(String name) {
+    private ApiOperation<StructDocument, StructDocument> getApiOperation(String name) {
         return operations.computeIfAbsent(name, operation -> {
             var shape = operationNames.get(name);
             if (shape == null) {

client/dynamic-client/src/main/java/software/amazon/smithy/java/dynamicclient/DynamicOperation.java

@@ -18,14 +18,14 @@
 import software.amazon.smithy.java.core.schema.TraitKey;
 import software.amazon.smithy.java.core.serde.TypeRegistry;
 import software.amazon.smithy.java.dynamicschemas.SchemaConverter;
-import software.amazon.smithy.java.dynamicschemas.WrappedDocument;
+import software.amazon.smithy.java.dynamicschemas.StructDocument;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.knowledge.ServiceIndex;
 import software.amazon.smithy.model.shapes.OperationShape;
 import software.amazon.smithy.model.shapes.ServiceShape;
 import software.amazon.smithy.model.shapes.ShapeId;
 
-public final class DynamicOperation implements ApiOperation<WrappedDocument, WrappedDocument> {
+public final class DynamicOperation implements ApiOperation<StructDocument, StructDocument> {
 
     private final ApiService service;
     private final Schema operationSchema;
@@ -84,12 +84,12 @@ public Schema outputSchema() {
     }
 
     @Override
-    public ShapeBuilder<WrappedDocument> inputBuilder() {
+    public ShapeBuilder<StructDocument> inputBuilder() {
         return SchemaConverter.createDocumentBuilder(inputSchema(), service.schema().id());
     }
 
     @Override
-    public ShapeBuilder<WrappedDocument> outputBuilder() {
+    public ShapeBuilder<StructDocument> outputBuilder() {
         return SchemaConverter.createDocumentBuilder(outputSchema(), service.schema().id());
     }
 

core/src/main/java/software/amazon/smithy/java/core/serde/document/DocumentUtils.java

@@ -48,7 +48,7 @@ public static void serializeNumber(ShapeSerializer serializer, Schema schema, Nu
         switch (schema.type()) {
             case BYTE -> serializer.writeByte(schema, value.byteValue());
             case SHORT -> serializer.writeShort(schema, value.shortValue());
-            case INTEGER -> serializer.writeInteger(schema, value.intValue());
+            case INTEGER, INT_ENUM -> serializer.writeInteger(schema, value.intValue());
             case LONG -> serializer.writeLong(schema, value.longValue());
             case FLOAT -> serializer.writeFloat(schema, value.floatValue());
             case DOUBLE -> serializer.writeDouble(schema, value.doubleValue());

core/src/main/java/software/amazon/smithy/java/core/serde/document/Documents.java

@@ -191,7 +191,7 @@ public void serializeContents(ShapeSerializer serializer) {
             serializer.writeList(schema, values, values.size(), (values, ser) -> {
                 for (var element : values) {
                     if (element == null) {
-                        ser.writeNull(PreludeSchemas.DOCUMENT);
+                        ser.writeNull(schema.listMember());
                     } else {
                         element.serialize(ser);
                     }

core/src/test/java/software/amazon/smithy/java/core/serde/document/ListDocumentTest.java

@@ -19,6 +19,7 @@
 import software.amazon.smithy.java.core.schema.Schema;
 import software.amazon.smithy.java.core.serde.ShapeSerializer;
 import software.amazon.smithy.java.core.serde.SpecificShapeSerializer;
+import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.shapes.ShapeType;
 
 public class ListDocumentTest {
@@ -106,13 +107,15 @@ public void writeDocument(Schema schema, Document value) {
 
                     @Override
                     public void writeString(Schema schema, String value) {
-                        assertThat(schema, equalTo(PreludeSchemas.STRING));
+                        // We don't attempt to change the schema of a list member. Since the given list member
+                        // was a document that didn't use a member schema, it's reflected here.
+                        assertThat(schema.id(), equalTo(ShapeId.from("smithy.api#String")));
                         writtenStrings.add(value);
                     }
 
                     @Override
                     public void writeNull(Schema schema) {
-                        assertThat(schema, equalTo(PreludeSchemas.DOCUMENT));
+                        assertThat(schema.id(), equalTo(ShapeId.from("smithy.api#Document$member")));
                         writtenStrings.add(null);
                     }
                 });

dynamic-schemas/src/main/java/software/amazon/smithy/java/dynamicschemas/ContentDocument.java

@@ -0,0 +1,187 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.java.dynamicschemas;
+
+import java.math.BigDecimal;
+import java.math.BigInteger;
+import java.nio.ByteBuffer;
+import java.time.Instant;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import software.amazon.smithy.java.core.schema.Schema;
+import software.amazon.smithy.java.core.schema.SerializableShape;
+import software.amazon.smithy.java.core.schema.ShapeBuilder;
+import software.amazon.smithy.java.core.serde.ShapeSerializer;
+import software.amazon.smithy.java.core.serde.document.Document;
+import software.amazon.smithy.model.shapes.ShapeType;
+
+/**
+ * A wrapper around another document that changes the serialized schema.
+ *
+ * @param document Document to wrap.
+ * @param schema Schema to use instead.
+ */
+record ContentDocument(Document document, Schema schema) implements Document {
+    @Override
+    public ShapeType type() {
+        return schema.type();
+    }
+
+    @Override
+    public void serialize(ShapeSerializer serializer) {
+        // If it is literally a document, then don't unwrap it.
+        if (type() == ShapeType.DOCUMENT) {
+            serializer.writeDocument(schema, this);
+        } else {
+            serializeContents(serializer);
+        }
+    }
+
+    @Override
+    public void serializeContents(ShapeSerializer serializer) {
+        switch (type()) {
+            case BOOLEAN -> serializer.writeBoolean(schema, asBoolean());
+            case BYTE -> serializer.writeByte(schema, asByte());
+            case SHORT -> serializer.writeShort(schema, asShort());
+            case INTEGER, INT_ENUM -> serializer.writeInteger(schema, asInteger());
+            case LONG -> serializer.writeLong(schema, asLong());
+            case FLOAT -> serializer.writeFloat(schema, asFloat());
+            case DOUBLE -> serializer.writeDouble(schema, asDouble());
+            case BIG_INTEGER -> serializer.writeBigInteger(schema, asBigInteger());
+            case BIG_DECIMAL -> serializer.writeBigDecimal(schema, asBigDecimal());
+            case STRING, ENUM -> serializer.writeString(schema, asString());
+            case TIMESTAMP -> serializer.writeTimestamp(schema, asTimestamp());
+            case BLOB -> serializer.writeBlob(schema, asBlob());
+            case DOCUMENT -> document.serializeContents(serializer);
+            case LIST, SET -> {
+                serializer.writeList(schema, asList(), size(), (values, ser) -> {
+                    for (var element : values) {
+                        if (element == null) {
+                            ser.writeNull(schema.listMember());
+                        } else {
+                            element.serialize(ser);
+                        }
+                    }
+                });
+            }
+            case MAP -> {
+                serializer.writeMap(schema, asStringMap(), size(), (members, s) -> {
+                    var key = schema.mapKeyMember();
+                    for (var entry : members.entrySet()) {
+                        if (entry.getValue() == null) {
+                            s.writeEntry(key, entry.getKey(), null, (t, v) -> v.writeNull(schema.mapValueMember()));
+                        } else {
+                            s.writeEntry(key, entry.getKey(), entry.getValue(), Document::serialize);
+                        }
+                    }
+                });
+            }
+            // Note that Structure and Union are always going to be a StructDocument and appear here.
+            default -> throw new UnsupportedOperationException("Unsupported type: " + type());
+        }
+    }
+
+    @Override
+    public BigDecimal asBigDecimal() {
+        return document.asBigDecimal();
+    }
+
+    @Override
+    public BigInteger asBigInteger() {
+        return document.asBigInteger();
+    }
+
+    @Override
+    public ByteBuffer asBlob() {
+        return document.asBlob();
+    }
+
+    @Override
+    public boolean asBoolean() {
+        return document.asBoolean();
+    }
+
+    @Override
+    public byte asByte() {
+        return document.asByte();
+    }
+
+    @Override
+    public double asDouble() {
+        return document.asDouble();
+    }
+
+    @Override
+    public float asFloat() {
+        return document.asFloat();
+    }
+
+    @Override
+    public int asInteger() {
+        return document.asInteger();
+    }
+
+    @Override
+    public List<Document> asList() {
+        return document.asList();
+    }
+
+    @Override
+    public long asLong() {
+        return document.asLong();
+    }
+
+    @Override
+    public Number asNumber() {
+        return document.asNumber();
+    }
+
+    @Override
+    public Object asObject() {
+        return document.asObject();
+    }
+
+    @Override
+    public <T extends SerializableShape> T asShape(ShapeBuilder<T> builder) {
+        return document.asShape(builder);
+    }
+
+    @Override
+    public short asShort() {
+        return document.asShort();
+    }
+
+    @Override
+    public String asString() {
+        return document.asString();
+    }
+
+    @Override
+    public Map<String, Document> asStringMap() {
+        return document.asStringMap();
+    }
+
+    @Override
+    public Instant asTimestamp() {
+        return document.asTimestamp();
+    }
+
+    @Override
+    public int size() {
+        return document.size();
+    }
+
+    @Override
+    public Document getMember(String memberName) {
+        return document.getMember(memberName);
+    }
+
+    @Override
+    public Set<String> getMemberNames() {
+        return document.getMemberNames();
+    }
+}