fix: s3 custom treatment getbucketlocation response (#897)

0marperez · web-flow · commit f62b7cfd7282 · 2023-08-08T11:35:36.000-04:00
diff --git a/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/core/RuntimeTypes.kt b/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/core/RuntimeTypes.kt
@@ -251,6 +251,7 @@ object RuntimeTypes {
             val XmlError = symbol("XmlError")
             val XmlSerializer = symbol("XmlSerializer")
             val XmlDeserializer = symbol("XmlDeserializer")
+            val XmlUnwrappedOutput = symbol("XmlUnwrappedOutput")
         }
 
         object SerdeFormUrl : RuntimeTypePackage(KotlinDependency.SERDE_FORM_URL) {
diff --git a/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/model/traits/UnwrappedXmlOutput.kt b/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/model/traits/UnwrappedXmlOutput.kt
@@ -0,0 +1,22 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package software.amazon.smithy.kotlin.codegen.model.traits
+
+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.AnnotationTrait
+
+/**
+ * Indicates the annotated structure is unwrapped XML output.
+ *
+ * Refer to this [s3 specific example](https://smithy.io/2.0/aws/customizations/s3-customizations.html#aws-customizations-s3unwrappedxmloutput-trait)
+ */
+class UnwrappedXmlOutput(node: ObjectNode) : AnnotationTrait(ID, node) {
+    companion object {
+        val ID: ShapeId = ShapeId.from("smithy.kotlin.traits#unwrappedXmlOutput")
+    }
+    constructor() : this(Node.objectNode())
+}
diff --git a/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/protocol/HttpBindingProtocolGenerator.kt b/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/protocol/HttpBindingProtocolGenerator.kt
@@ -12,7 +12,10 @@ import software.amazon.smithy.kotlin.codegen.core.*
 import software.amazon.smithy.kotlin.codegen.lang.KotlinTypes
 import software.amazon.smithy.kotlin.codegen.lang.toEscapedLiteral
 import software.amazon.smithy.kotlin.codegen.model.*
-import software.amazon.smithy.kotlin.codegen.rendering.serde.*
+import software.amazon.smithy.kotlin.codegen.rendering.serde.deserializerName
+import software.amazon.smithy.kotlin.codegen.rendering.serde.formatInstant
+import software.amazon.smithy.kotlin.codegen.rendering.serde.parseInstant
+import software.amazon.smithy.kotlin.codegen.rendering.serde.serializerName
 import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.knowledge.HttpBinding
 import software.amazon.smithy.model.shapes.*
diff --git a/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/protocol/HttpProtocolClientGenerator.kt b/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/protocol/HttpProtocolClientGenerator.kt
@@ -9,8 +9,11 @@ import software.amazon.smithy.kotlin.codegen.core.*
 import software.amazon.smithy.kotlin.codegen.integration.SectionId
 import software.amazon.smithy.kotlin.codegen.integration.SectionKey
 import software.amazon.smithy.kotlin.codegen.lang.KotlinTypes
-import software.amazon.smithy.kotlin.codegen.model.*
+import software.amazon.smithy.kotlin.codegen.model.getTrait
+import software.amazon.smithy.kotlin.codegen.model.hasStreamingMember
+import software.amazon.smithy.kotlin.codegen.model.hasTrait
 import software.amazon.smithy.kotlin.codegen.model.knowledge.AuthIndex
+import software.amazon.smithy.kotlin.codegen.model.operationSignature
 import software.amazon.smithy.kotlin.codegen.rendering.auth.AuthSchemeProviderAdapterGenerator
 import software.amazon.smithy.kotlin.codegen.rendering.auth.IdentityProviderConfigGenerator
 import software.amazon.smithy.kotlin.codegen.rendering.endpoints.EndpointResolverAdapterGenerator
@@ -36,11 +39,6 @@ abstract class HttpProtocolClientGenerator(
         val OperationShape: SectionKey<OperationShape> = SectionKey("OperationShape")
     }
 
-    object OperationDeserializerBinding : SectionId {
-        // Context for operation being codegened at the time of section invocation
-        val Operation: SectionKey<OperationShape> = SectionKey("Operation")
-    }
-
     object OperationTelemetryBuilder : SectionId {
         val Operation: SectionKey<OperationShape> = SectionKey("Operation")
     }
@@ -204,12 +202,10 @@ abstract class HttpProtocolClientGenerator(
                 }
             }
 
-            writer.declareSection(OperationDeserializerBinding, mapOf(OperationDeserializerBinding.Operation to op)) {
-                if (outputShape.isPresent) {
-                    write("deserializer = ${op.deserializerName()}()")
-                } else {
-                    write("deserializer = UnitDeserializer")
-                }
+            if (outputShape.isPresent) {
+                writer.write("deserializer = ${op.deserializerName()}()")
+            } else {
+                writer.write("deserializer = UnitDeserializer")
             }
 
             // execution context
diff --git a/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/serde/XmlSerdeDescriptorGenerator.kt b/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/serde/XmlSerdeDescriptorGenerator.kt
@@ -6,9 +6,16 @@
 package software.amazon.smithy.kotlin.codegen.rendering.serde
 
 import software.amazon.smithy.codegen.core.Symbol
-import software.amazon.smithy.kotlin.codegen.core.*
-import software.amazon.smithy.kotlin.codegen.model.*
+import software.amazon.smithy.kotlin.codegen.core.RenderingContext
+import software.amazon.smithy.kotlin.codegen.core.RuntimeTypes
+import software.amazon.smithy.kotlin.codegen.core.addImport
+import software.amazon.smithy.kotlin.codegen.core.defaultName
+import software.amazon.smithy.kotlin.codegen.model.expectShape
+import software.amazon.smithy.kotlin.codegen.model.expectTrait
+import software.amazon.smithy.kotlin.codegen.model.getTrait
+import software.amazon.smithy.kotlin.codegen.model.hasTrait
 import software.amazon.smithy.kotlin.codegen.model.traits.SyntheticClone
+import software.amazon.smithy.kotlin.codegen.model.traits.UnwrappedXmlOutput
 import software.amazon.smithy.kotlin.codegen.utils.dq
 import software.amazon.smithy.model.shapes.*
 import software.amazon.smithy.model.traits.*
@@ -50,6 +57,10 @@ open class XmlSerdeDescriptorGenerator(
             objTraits.add(serdeTrait)
         }
 
+        if (objectShape.hasTrait<UnwrappedXmlOutput>()) {
+            objTraits.add(SerdeXml.XmlUnwrappedOutput)
+        }
+
         return objTraits
     }
 
diff --git a/runtime/serde/serde-xml/common/src/aws/smithy/kotlin/runtime/serde/xml/XmlDeserializer.kt b/runtime/serde/serde-xml/common/src/aws/smithy/kotlin/runtime/serde/xml/XmlDeserializer.kt
@@ -11,6 +11,8 @@ import aws.smithy.kotlin.runtime.content.BigInteger
 import aws.smithy.kotlin.runtime.content.Document
 import aws.smithy.kotlin.runtime.serde.*
 
+private const val FIRST_FIELD_INDEX: Int = 0
+
 // Represents aspects of SdkFieldDescriptor that are particular to the Xml format
 internal sealed class FieldLocation {
     // specifies the mapping to a sdk field index
@@ -69,7 +71,8 @@ public class XmlDeserializer(
             throw DeserializationException("Expected last parsed token to be ${XmlToken.BeginElement::class} but was ${reader.lastToken}")
         }
 
-        return XmlStructDeserializer(descriptor, reader.subTreeReader(), parentToken, attribFields)
+        val unwrapped = descriptor.hasTrait<XmlUnwrappedOutput>()
+        return XmlStructDeserializer(descriptor, reader.subTreeReader(XmlStreamReader.SubtreeStartDepth.CURRENT), parentToken, attribFields, unwrapped)
     }
 
     override fun deserializeList(descriptor: SdkFieldDescriptor): Deserializer.ElementIterator {
@@ -213,16 +216,22 @@ internal class XmlListDeserializer(
  * @param parentToken initial token of associated structure
  * @param parsedFieldLocations list of [FieldLocation] representing values able to be loaded into deserialized instances
  */
-internal class XmlStructDeserializer(
+private class XmlStructDeserializer(
     private val objDescriptor: SdkObjectDescriptor,
-    private val reader: XmlStreamReader,
+    reader: XmlStreamReader,
     private val parentToken: XmlToken.BeginElement,
     private val parsedFieldLocations: MutableList<FieldLocation> = mutableListOf(),
+    private val unwrapped: Boolean,
 ) : Deserializer.FieldIterator {
     // Used to track direct deserialization or further nesting between calls to findNextFieldIndex() and deserialize<Type>()
     private var reentryFlag: Boolean = false
 
+    private val reader: XmlStreamReader = if (unwrapped) reader else reader.subTreeReader(XmlStreamReader.SubtreeStartDepth.CHILD)
+
     override fun findNextFieldIndex(): Int? {
+        if (unwrapped) {
+            return if (reader.peek() is XmlToken.Text) FIRST_FIELD_INDEX else null
+        }
         if (inNestedMode()) {
             // Returning from a nested struct call.  Nested deserializer consumed
             // tokens so clear them here to avoid processing stale state
@@ -251,6 +260,10 @@ internal class XmlStructDeserializer(
     }
 
     private fun <T> deserializeValue(transform: ((String) -> T)): T {
+        if (unwrapped) {
+            val value = reader.takeNextAs<XmlToken.Text>().value ?: ""
+            return transform(value)
+        }
         // Set and validate mode
         reentryFlag = false
         if (parsedFieldLocations.isEmpty()) throw DeserializationException("matchedFields is empty, was findNextFieldIndex() called?")
@@ -383,16 +396,20 @@ private fun SdkObjectDescriptor.fieldTokenMatcher(fieldDescriptor: SdkFieldDescr
     return fieldDescriptor.nameMatches(beginElement.name.tag)
 }
 
-// Return the next token of the specified type or throw [DeserializerStateException] if incorrect type.
+/**
+ * Return the next token of the specified type or throw [DeserializationException] if incorrect type.
+ */
 internal inline fun <reified TExpected : XmlToken> XmlStreamReader.takeNextAs(): TExpected {
     val token = this.nextToken() ?: throw DeserializationException("Expected ${TExpected::class} but instead found null")
     requireToken<TExpected>(token)
     return token as TExpected
 }
 
-// require that the given token be of type [TExpected] or else throw an exception
+/**
+ * Require that the given token be of type [TExpected] or else throw an exception
+ */
 internal inline fun <reified TExpected> requireToken(token: XmlToken) {
     if (token::class != TExpected::class) {
-        throw DeserializationException("expected ${TExpected::class}; found ${token::class} ($token)")
+        throw DeserializationException("Expected ${TExpected::class}; found ${token::class} ($token)")
     }
 }
diff --git a/runtime/serde/serde-xml/common/src/aws/smithy/kotlin/runtime/serde/xml/XmlFieldTraits.kt b/runtime/serde/serde-xml/common/src/aws/smithy/kotlin/runtime/serde/xml/XmlFieldTraits.kt
@@ -67,6 +67,14 @@ public data class XmlCollectionName(
 @InternalApi
 public object Flattened : FieldTrait
 
+/**
+ * Specifies that an object is XML unwrapped response.
+ *
+ * Refer to: [s3 specific example](https://smithy.io/2.0/aws/customizations/s3-customizations.html#aws-customizations-s3unwrappedxmloutput-trait)
+ */
+@InternalApi
+public object XmlUnwrappedOutput : FieldTrait
+
 /**
  * Denotes a structure that represents an error.  There are special rules for error deserialization
  * in various XML-based protocols. This trait provides necessary context to the deserializer to properly
diff --git a/runtime/serde/serde-xml/common/test/aws/smithy/kotlin/runtime/serde/xml/XmlDeserializerStructTest.kt b/runtime/serde/serde-xml/common/test/aws/smithy/kotlin/runtime/serde/xml/XmlDeserializerStructTest.kt
@@ -200,6 +200,57 @@ class XmlDeserializerStructTest {
         println(bst.nested?.nested)
     }
 
+    class BasicUnwrappedStructTest {
+        var x: String? = null
+
+        companion object {
+            val X_DESCRIPTOR = SdkFieldDescriptor(SerialKind.String, XmlSerialName("x"))
+            private val OBJ_DESCRIPTOR = SdkObjectDescriptor.build {
+                trait(XmlSerialName("payload"))
+                trait(XmlUnwrappedOutput)
+                field(X_DESCRIPTOR)
+            }
+
+            fun deserialize(deserializer: Deserializer): BasicUnwrappedStructTest {
+                val result = BasicUnwrappedStructTest()
+                deserializer.deserializeStruct(OBJ_DESCRIPTOR) {
+                    loop@ while (true) {
+                        when (findNextFieldIndex()) {
+                            X_DESCRIPTOR.index -> result.x = deserializeString()
+                            null -> break@loop
+                            else -> throw DeserializationException(IllegalStateException("unexpected field in BasicUnwrappedStructTest deserializer"))
+                        }
+                    }
+                }
+                return result
+            }
+        }
+    }
+
+    @Test
+    fun itHandlesBasicUnwrappedStructs() {
+        val payload = """
+            <x>text</x>
+        """.encodeToByteArray()
+
+        val deserializer = XmlDeserializer(payload)
+        val bst = BasicUnwrappedStructTest.deserialize(deserializer)
+
+        assertEquals("text", bst.x)
+    }
+
+    @Test
+    fun itHandlesBasicUnwrappedStructsWithNullValues() {
+        val payload = """
+            <x></x>
+        """.encodeToByteArray()
+
+        val deserializer = XmlDeserializer(payload)
+        val bst = BasicUnwrappedStructTest.deserialize(deserializer)
+
+        assertEquals(null, bst.x)
+    }
+
     class AliasStruct {
         var message: String? = null
         var attribute: String? = null

Original file line number	Diff line number	Diff line change
`@@ -251,6 +251,7 @@ object RuntimeTypes {`
`251`	`251`	`val XmlError = symbol("XmlError")`
`252`	`252`	`val XmlSerializer = symbol("XmlSerializer")`
`253`	`253`	`val XmlDeserializer = symbol("XmlDeserializer")`
	`254`	`+ val XmlUnwrappedOutput = symbol("XmlUnwrappedOutput")`
`254`	`255`	`}`
`255`	`256`
`256`	`257`	`object SerdeFormUrl : RuntimeTypePackage(KotlinDependency.SERDE_FORM_URL) {`