fix: Allow customization of HTTP body members per-protocol (#907)

jbelkins · web-flow · commit b5b67b33e150 · 2025-02-17T10:04:00.000-06:00
diff --git a/smithy-swift-codegen/src/main/kotlin/software/amazon/smithy/swift/codegen/integration/HTTPBindingProtocolGenerator.kt b/smithy-swift-codegen/src/main/kotlin/software/amazon/smithy/swift/codegen/integration/HTTPBindingProtocolGenerator.kt
@@ -68,22 +68,20 @@ import software.amazon.smithy.swift.codegen.model.hasEventStreamMember
 import software.amazon.smithy.swift.codegen.model.hasTrait
 import software.amazon.smithy.swift.codegen.model.isInputEventStream
 import software.amazon.smithy.swift.codegen.model.isOutputEventStream
-import software.amazon.smithy.swift.codegen.model.targetOrSelf
 import software.amazon.smithy.swift.codegen.supportsStreamingAndIsRPC
 import software.amazon.smithy.swift.codegen.swiftmodules.ClientRuntimeTypes
 import software.amazon.smithy.swift.codegen.utils.ModelFileUtils
 import software.amazon.smithy.utils.OptionalUtils
 import java.util.Optional
 import java.util.logging.Logger
 
-private val Shape.isStreaming: Boolean
+val Shape.isEventStreaming: Boolean
     get() = hasTrait<StreamingTrait>() && isUnionShape
 
 /**
  * Checks to see if shape is in the body of the http request
  */
 fun Shape.isInHttpBody(): Boolean {
-    // TODO fix the edge case: a shape which is an operational input (i.e. has members bound to HTTP semantics) could be re-used elsewhere not as an operation input which means everything is in the body
     val hasNoHttpTraitsOutsideOfPayload =
         !this.hasTrait(HttpLabelTrait::class.java) &&
             !this.hasTrait(HttpHeaderTrait::class.java) &&
@@ -171,19 +169,7 @@ abstract class HTTPBindingProtocolGenerator(
                     .definitionFile(filename)
                     .name(symbolName)
                     .build()
-            var httpBodyMembers =
-                shape
-                    .members()
-                    .filter { it.isInHttpBody() }
-                    .toList()
-            if (supportsStreamingAndIsRPC(ctx.protocol)) {
-                // For RPC protocols that support event streaming, we need to send initial request
-                // with streaming member excluded during encoding the input struct.
-                httpBodyMembers =
-                    httpBodyMembers.filter {
-                        !it.targetOrSelf(ctx.model).isStreaming
-                    }
-            }
+            val httpBodyMembers = httpBodyMembers(ctx, shape)
             if (httpBodyMembers.isNotEmpty() || shouldRenderEncodableConformance) {
                 ctx.delegator.useShapeWriter(encodeSymbol) { writer ->
                     writer.openBlock(
@@ -208,7 +194,7 @@ abstract class HTTPBindingProtocolGenerator(
     override fun generateCodableConformanceForNestedTypes(ctx: ProtocolGenerator.GenerationContext) {
         val nestedShapes =
             resolveShapesNeedingCodableConformance(ctx)
-                .filter { !it.isStreaming }
+                .filter { !it.isEventStreaming }
         for (shape in nestedShapes) {
             renderCodableExtension(ctx, shape)
         }
@@ -567,6 +553,11 @@ abstract class HTTPBindingProtocolGenerator(
             messageUnmarshallableGenerator.render(streamingMember)
         }
     }
+
+    abstract fun httpBodyMembers(
+        ctx: ProtocolGenerator.GenerationContext,
+        shape: Shape,
+    ): List<MemberShape>
 }
 
 class DefaultServiceConfig(
diff --git a/smithy-swift-codegen/src/main/kotlin/software/amazon/smithy/swift/codegen/protocols/rpcv2cbor/RpcV2CborProtocolGenerator.kt b/smithy-swift-codegen/src/main/kotlin/software/amazon/smithy/swift/codegen/protocols/rpcv2cbor/RpcV2CborProtocolGenerator.kt
@@ -1,7 +1,9 @@
 package software.amazon.smithy.swift.codegen.protocols.rpcv2cbor
 
 import software.amazon.smithy.codegen.core.Symbol
+import software.amazon.smithy.model.shapes.MemberShape
 import software.amazon.smithy.model.shapes.OperationShape
+import software.amazon.smithy.model.shapes.Shape
 import software.amazon.smithy.model.shapes.ShapeId
 import software.amazon.smithy.model.traits.StreamingTrait
 import software.amazon.smithy.model.traits.UnitTypeTrait
@@ -89,6 +91,14 @@ class RpcV2CborProtocolGenerator(
         }
     }
 
+    override fun httpBodyMembers(
+        ctx: ProtocolGenerator.GenerationContext,
+        shape: Shape,
+    ): List<MemberShape> =
+        shape
+            .members()
+            .toList()
+
     /**
      * @return whether the operation input does _not_ target the unit shape ([UnitTypeTrait.UNIT])
      */
diff --git a/smithy-swift-codegen/src/test/kotlin/software/amazon/smithy/swift/codegen/protocolgeneratormocks/MockHTTPAWSJson11ProtocolGenerator.kt b/smithy-swift-codegen/src/test/kotlin/software/amazon/smithy/swift/codegen/protocolgeneratormocks/MockHTTPAWSJson11ProtocolGenerator.kt
@@ -7,7 +7,9 @@ package software.amazon.smithy.swift.codegen.protocolgeneratormocks
 
 import software.amazon.smithy.aws.traits.protocols.AwsJson1_1Trait
 import software.amazon.smithy.model.pattern.UriPattern
+import software.amazon.smithy.model.shapes.MemberShape
 import software.amazon.smithy.model.shapes.OperationShape
+import software.amazon.smithy.model.shapes.Shape
 import software.amazon.smithy.model.shapes.ShapeId
 import software.amazon.smithy.model.traits.HttpTrait
 import software.amazon.smithy.swift.codegen.SwiftWriter
@@ -19,7 +21,9 @@ import software.amazon.smithy.swift.codegen.integration.HttpProtocolUnitTestErro
 import software.amazon.smithy.swift.codegen.integration.HttpProtocolUnitTestRequestGenerator
 import software.amazon.smithy.swift.codegen.integration.HttpProtocolUnitTestResponseGenerator
 import software.amazon.smithy.swift.codegen.integration.ProtocolGenerator
+import software.amazon.smithy.swift.codegen.integration.isEventStreaming
 import software.amazon.smithy.swift.codegen.integration.protocols.core.StaticHttpBindingResolver
+import software.amazon.smithy.swift.codegen.model.targetOrSelf
 import software.amazon.smithy.swift.codegen.requestandresponse.TestHttpProtocolClientGeneratorFactory
 
 class MockJsonHttpBindingResolver(
@@ -87,4 +91,15 @@ class MockHTTPAWSJson11ProtocolGenerator : HTTPBindingProtocolGenerator(MockAWSJ
         ctx: ProtocolGenerator.GenerationContext,
         defaultContentType: String,
     ): HttpBindingResolver = MockJsonHttpBindingResolver(ctx, defaultContentType)
+
+    override fun httpBodyMembers(
+        ctx: ProtocolGenerator.GenerationContext,
+        shape: Shape,
+    ): List<MemberShape> =
+        shape
+            .members()
+            // For RPC protocols that support event streaming, we need to send initial request
+            // with streaming member excluded during encoding the input struct.
+            .filter { !it.targetOrSelf(ctx.model).isEventStreaming }
+            .toList()
 }
diff --git a/smithy-swift-codegen/src/test/kotlin/software/amazon/smithy/swift/codegen/protocolgeneratormocks/MockHTTPEC2QueryProtocolGenerator.kt b/smithy-swift-codegen/src/test/kotlin/software/amazon/smithy/swift/codegen/protocolgeneratormocks/MockHTTPEC2QueryProtocolGenerator.kt
@@ -7,7 +7,9 @@ package software.amazon.smithy.swift.codegen.protocolgeneratormocks
 
 import software.amazon.smithy.aws.traits.protocols.Ec2QueryTrait
 import software.amazon.smithy.model.pattern.UriPattern
+import software.amazon.smithy.model.shapes.MemberShape
 import software.amazon.smithy.model.shapes.OperationShape
+import software.amazon.smithy.model.shapes.Shape
 import software.amazon.smithy.model.shapes.ShapeId
 import software.amazon.smithy.model.traits.HttpTrait
 import software.amazon.smithy.swift.codegen.integration.DefaultHTTPProtocolCustomizations
@@ -77,4 +79,12 @@ class MockHTTPEC2QueryProtocolGenerator : HTTPBindingProtocolGenerator(MockEC2Qu
             getProtocolHttpBindingResolver(ctx, defaultContentType),
         ).generateProtocolTests()
     }
+
+    override fun httpBodyMembers(
+        ctx: ProtocolGenerator.GenerationContext,
+        shape: Shape,
+    ): List<MemberShape> =
+        shape
+            .members()
+            .toList()
 }
diff --git a/smithy-swift-codegen/src/test/kotlin/software/amazon/smithy/swift/codegen/protocolgeneratormocks/MockHTTPRPCv2CBORProtocolGenerator.kt b/smithy-swift-codegen/src/test/kotlin/software/amazon/smithy/swift/codegen/protocolgeneratormocks/MockHTTPRPCv2CBORProtocolGenerator.kt
@@ -5,7 +5,9 @@ package software.amazon.smithy.swift.codegen.protocolgeneratormocks
  * SPDX-License-Identifier: Apache-2.0.
  */
 
+import software.amazon.smithy.model.shapes.MemberShape
 import software.amazon.smithy.model.shapes.OperationShape
+import software.amazon.smithy.model.shapes.Shape
 import software.amazon.smithy.model.shapes.ShapeId
 import software.amazon.smithy.protocol.traits.Rpcv2CborTrait
 import software.amazon.smithy.swift.codegen.integration.DefaultHTTPProtocolCustomizations
@@ -53,4 +55,12 @@ class MockHTTPRPCv2CBORProtocolGenerator : HTTPBindingProtocolGenerator(MockRPCv
             getProtocolHttpBindingResolver(ctx, defaultContentType),
         ).generateProtocolTests()
     }
+
+    override fun httpBodyMembers(
+        ctx: ProtocolGenerator.GenerationContext,
+        shape: Shape,
+    ): List<MemberShape> =
+        shape
+            .members()
+            .toList()
 }
diff --git a/smithy-swift-codegen/src/test/kotlin/software/amazon/smithy/swift/codegen/protocolgeneratormocks/MockHTTPRestJsonProtocolGenerator.kt b/smithy-swift-codegen/src/test/kotlin/software/amazon/smithy/swift/codegen/protocolgeneratormocks/MockHTTPRestJsonProtocolGenerator.kt
@@ -6,7 +6,9 @@ package software.amazon.smithy.swift.codegen.protocolgeneratormocks
  */
 
 import software.amazon.smithy.aws.traits.protocols.RestJson1Trait
+import software.amazon.smithy.model.shapes.MemberShape
 import software.amazon.smithy.model.shapes.OperationShape
+import software.amazon.smithy.model.shapes.Shape
 import software.amazon.smithy.model.shapes.ShapeId
 import software.amazon.smithy.swift.codegen.integration.DefaultHTTPProtocolCustomizations
 import software.amazon.smithy.swift.codegen.integration.HTTPBindingProtocolGenerator
@@ -15,6 +17,7 @@ import software.amazon.smithy.swift.codegen.integration.HttpProtocolUnitTestErro
 import software.amazon.smithy.swift.codegen.integration.HttpProtocolUnitTestRequestGenerator
 import software.amazon.smithy.swift.codegen.integration.HttpProtocolUnitTestResponseGenerator
 import software.amazon.smithy.swift.codegen.integration.ProtocolGenerator
+import software.amazon.smithy.swift.codegen.integration.isInHttpBody
 import software.amazon.smithy.swift.codegen.requestandresponse.TestHttpProtocolClientGeneratorFactory
 
 class MockRestJsonHTTPProtocolCustomizations : DefaultHTTPProtocolCustomizations()
@@ -53,4 +56,13 @@ class MockHTTPRestJsonProtocolGenerator : HTTPBindingProtocolGenerator(MockRestJ
             getProtocolHttpBindingResolver(ctx, defaultContentType),
         ).generateProtocolTests()
     }
+
+    override fun httpBodyMembers(
+        ctx: ProtocolGenerator.GenerationContext,
+        shape: Shape,
+    ): List<MemberShape> =
+        shape
+            .members()
+            .filter { it.isInHttpBody() }
+            .toList()
 }
diff --git a/smithy-swift-codegen/src/test/kotlin/software/amazon/smithy/swift/codegen/protocolgeneratormocks/MockHTTPRestXMLProtocolGenerator.kt b/smithy-swift-codegen/src/test/kotlin/software/amazon/smithy/swift/codegen/protocolgeneratormocks/MockHTTPRestXMLProtocolGenerator.kt
@@ -6,7 +6,9 @@ package software.amazon.smithy.swift.codegen.protocolgeneratormocks
  */
 
 import software.amazon.smithy.aws.traits.protocols.RestXmlTrait
+import software.amazon.smithy.model.shapes.MemberShape
 import software.amazon.smithy.model.shapes.OperationShape
+import software.amazon.smithy.model.shapes.Shape
 import software.amazon.smithy.model.shapes.ShapeId
 import software.amazon.smithy.swift.codegen.integration.DefaultHTTPProtocolCustomizations
 import software.amazon.smithy.swift.codegen.integration.HTTPBindingProtocolGenerator
@@ -15,6 +17,7 @@ import software.amazon.smithy.swift.codegen.integration.HttpProtocolUnitTestErro
 import software.amazon.smithy.swift.codegen.integration.HttpProtocolUnitTestRequestGenerator
 import software.amazon.smithy.swift.codegen.integration.HttpProtocolUnitTestResponseGenerator
 import software.amazon.smithy.swift.codegen.integration.ProtocolGenerator
+import software.amazon.smithy.swift.codegen.integration.isInHttpBody
 import software.amazon.smithy.swift.codegen.requestandresponse.TestHttpProtocolClientGeneratorFactory
 
 class MockRestXMLHTTPProtocolCustomizations : DefaultHTTPProtocolCustomizations()
@@ -53,4 +56,13 @@ class MockHTTPRestXMLProtocolGenerator : HTTPBindingProtocolGenerator(MockRestXM
             getProtocolHttpBindingResolver(ctx, defaultContentType),
         ).generateProtocolTests()
     }
+
+    override fun httpBodyMembers(
+        ctx: ProtocolGenerator.GenerationContext,
+        shape: Shape,
+    ): List<MemberShape> =
+        shape
+            .members()
+            .filter { it.isInHttpBody() }
+            .toList()
 }
