codegen: Add support for REST-XML Union responses (#993)

Adds support for deserializing REST-XML union resources Fixes #939 Related to smithy-lang/smithy#678
aws · Dec 22, 2020 · b58dcf8 · b58dcf8
1 parent 4947df4
commit b58dcf8
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 6 deletions.
diff --git a/...go-codegen/src/main/java/software/amazon/smithy/aws/go/codegen/JsonShapeDeserVisitor.java b/...go-codegen/src/main/java/software/amazon/smithy/aws/go/codegen/JsonShapeDeserVisitor.java
@@ -244,8 +244,6 @@ protected void deserializeUnion(GenerationContext context, UnionShape shape) {
                             symbol.getNamespace()
                     ).build();
                     writer.write("uv = &$T{Tag: key}", unknownMemberSymbol);
-                    writer.writeDocs("TODO: FIX ME");
-                    writer.write("_ = value");
                     writer.write("break loop");
                 });
             });

diff --git a/...-go-codegen/src/main/java/software/amazon/smithy/aws/go/codegen/XmlShapeDeserVisitor.java b/...-go-codegen/src/main/java/software/amazon/smithy/aws/go/codegen/XmlShapeDeserVisitor.java
@@ -11,6 +11,8 @@
 import software.amazon.smithy.codegen.core.SymbolProvider;
 import software.amazon.smithy.go.codegen.GoWriter;
 import software.amazon.smithy.go.codegen.SmithyGoDependency;
+import software.amazon.smithy.go.codegen.SymbolUtils;
+import software.amazon.smithy.go.codegen.UnionGenerator;
 import software.amazon.smithy.go.codegen.integration.DocumentShapeDeserVisitor;
 import software.amazon.smithy.go.codegen.integration.ProtocolGenerator;
 import software.amazon.smithy.go.codegen.integration.ProtocolGenerator.GenerationContext;
@@ -404,10 +406,67 @@ protected void deserializeDocument(GenerationContext context, DocumentShape shap
     @Override
     protected void deserializeUnion(GenerationContext context, UnionShape shape) {
         GoWriter writer = context.getWriter();
-        // TODO: implement union deserialization
-        // The tricky part is going to be accumulating bytes for unknown members.
-        LOGGER.warning("Union type is currently unsupported for XML deserialization.");
-        context.getWriter().writeDocs("TODO: implement union serialization.");
+        SymbolProvider symbolProvider = context.getSymbolProvider();
+        Symbol symbol = symbolProvider.toSymbol(shape);
+        Model model = context.getModel();
+
+        writer.write("var uv $P", symbol);
+        writer.write("var memberFound bool");
+
+        // Iterate through the decoder. The member visitor will handle popping
+        // xml tokens enclosed within a xml start and end element.
+        writer.openBlock("for {", "}", () -> {
+            writer.write("t, done, err := decoder.Token()");
+            writer.write("if err != nil { return err }");
+            writer.write("if done { break }");
+            writer.openBlock("if memberFound {", "}", () -> {
+                writer.write("if err = decoder.Decoder.Skip(); err != nil { return err }");
+            });
+
+            // Create a new decoder for each member
+            writer.write("originalDecoder := decoder");
+            writer.write("decoder = smithyxml.WrapNodeDecoder(originalDecoder.Decoder, t)");
+            writer.insertTrailingNewline();
+
+            writer.openBlock("switch {", "}", () -> {
+                Set<MemberShape> members = new TreeSet<>(shape.members());
+                for (MemberShape member : members) {
+                    // check if member is not a document binding or has a xmlAttribute trait
+                    if (!memberFilter.test(member) || member.hasTrait(XmlAttributeTrait.ID)) {
+                        continue;
+                    }
+                    Symbol targetSymbol = symbolProvider.toSymbol(member);
+                    Symbol memberSymbol = SymbolUtils.createValueSymbolBuilder(
+                            symbolProvider.toMemberName(member),
+                            symbol.getNamespace()
+                    ).build();
+
+                    writer.addUseImports(SmithyGoDependency.STRINGS);
+
+                    String serializedMemberName = getSerializedMemberName(member);
+                    writer.openBlock("case strings.EqualFold($S, t.Name.Local):", "", serializedMemberName, () -> {
+                        writer.write("var mv $P", targetSymbol);
+                        model.expectShape(member.getTarget()).accept(getMemberDeserVisitor(member, "mv", false));
+                        writer.write("uv = &$T{Value: mv}", memberSymbol);
+                        writer.write("memberFound = true");
+                    });
+                }
+
+                writer.openBlock("default:", "", () -> {
+                    // This is the function to take a value and convert it to the union type.
+                    Symbol unknownMemberSymbol = SymbolUtils.createValueSymbolBuilder(
+                            UnionGenerator.UNKNOWN_MEMBER_NAME,
+                            symbol.getNamespace()
+                    ).build();
+                    writer.write("uv = &$T{Tag: t.Name.Local}", unknownMemberSymbol);
+                    writer.write("memberFound = true");
+                });
+            });
+            // re-assign the  original decoder
+            writer.write("decoder = originalDecoder");
+        });
+
+        writer.write("*v = uv");
         writer.write("return nil");
     }
 }