feat: support initial-request and initial-response event stream types for RPC-bound protocols (#1026)

Author: lauzadis

.changes/d6c229e5-12a0-4d32-97d9-a7368a4d71dd.json

@@ -0,0 +1,5 @@
+{
+    "id": "d6c229e5-12a0-4d32-97d9-a7368a4d71dd",
+    "type": "feature",
+    "description": "Support initial-request and initial-response for event streams using RPC-based protocols"
+}

codegen/smithy-aws-kotlin-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/customization/RemoveEventStreamOperations.kt

@@ -15,9 +15,27 @@ import software.amazon.smithy.kotlin.codegen.rendering.protocol.ProtocolGenerato
 import software.amazon.smithy.kotlin.codegen.rendering.serde.StructuredDataParserGenerator
 import software.amazon.smithy.kotlin.codegen.rendering.serde.bodyDeserializer
 import software.amazon.smithy.kotlin.codegen.rendering.serde.bodyDeserializerName
+import software.amazon.smithy.kotlin.codegen.utils.getOrNull
 import software.amazon.smithy.model.shapes.*
 import software.amazon.smithy.model.traits.EventHeaderTrait
 import software.amazon.smithy.model.traits.EventPayloadTrait
+import software.amazon.smithy.model.traits.StreamingTrait
+
+/**
+ * A set of RPC-bound Smithy protocols
+ */
+val RPC_BOUND_PROTOCOLS = setOf(
+    "awsJson1_0",
+    "awsJson1_1",
+    "awsQuery",
+    "ec2Query",
+)
+
+/**
+ * Represents whether the given ShapeId represents an RPC-bound Smithy protocol
+ */
+internal val ShapeId.isRpcBoundProtocol: Boolean
+    get() = RPC_BOUND_PROTOCOLS.contains(name)
 
 /**
  * Implements rendering deserialize implementation for event streams implemented using the
@@ -60,26 +78,29 @@ class EventStreamParserGenerator(
         val streamShape = ctx.model.expectShape<UnionShape>(streamingMember.target)
         val streamSymbol = ctx.symbolProvider.toSymbol(streamShape)
 
-        // TODO - handle RPC bound protocol bindings where the initial response is bound to an event stream document
-        //        possibly by decoding the first Message
-
         val messageTypeSymbol = RuntimeTypes.AwsEventStream.MessageType
         val baseExceptionSymbol = ExceptionBaseClassGenerator.baseExceptionSymbol(ctx.settings)
 
         writer.write("val chan = body.#T() ?: return", RuntimeTypes.Http.toSdkByteReadChannel)
-        writer.write("val events = #T(chan)", RuntimeTypes.AwsEventStream.decodeFrames)
-            .indent()
+        writer.write("val frames = #T(chan)", RuntimeTypes.AwsEventStream.decodeFrames)
+        if (ctx.protocol.isRpcBoundProtocol) {
+            renderDeserializeInitialResponse(ctx, output, writer)
+        } else {
+            writer.write("val events = frames")
+        }
+        writer.indent()
             .withBlock(".#T { message ->", "}", RuntimeTypes.KotlinxCoroutines.Flow.map) {
-                withBlock("when(val mt = message.#T()) {", "}", RuntimeTypes.AwsEventStream.MessageTypeExt) {
-                    withBlock("is #T.Event -> when(mt.shapeType) {", "}", messageTypeSymbol) {
+                withBlock("when (val mt = message.#T()) {", "}", RuntimeTypes.AwsEventStream.MessageTypeExt) {
+                    withBlock("is #T.Event -> when (mt.shapeType) {", "}", messageTypeSymbol) {
                         streamShape.filterEventStreamErrors(ctx.model).forEach { member ->
                             withBlock("#S -> {", "}", member.memberName) {
                                 renderDeserializeEventVariant(ctx, streamSymbol, member, writer)
                             }
                         }
+
                         write("else -> #T.SdkUnknown", streamSymbol)
                     }
-                    withBlock("is #T.Exception -> when(mt.shapeType){", "}", messageTypeSymbol) {
+                    withBlock("is #T.Exception -> when (mt.shapeType) {", "}", messageTypeSymbol) {
                         // errors are completely bound to payload (at least according to design docs)
                         val errorMembers = streamShape.members().filter {
                             val target = ctx.model.expectShape(it.target)
@@ -101,7 +122,9 @@ class EventStreamParserGenerator(
                 }
             }
             .dedent()
-            .write("builder.#L = events", streamingMember.defaultName())
+            .write("")
+
+        writer.write("builder.#L = events", streamingMember.defaultName())
     }
 
     private fun renderDeserializeEventVariant(ctx: ProtocolGenerator.GenerationContext, unionSymbol: Symbol, member: MemberShape, writer: KotlinWriter) {
@@ -167,6 +190,48 @@ class EventStreamParserGenerator(
         writer.write("#T.#L(e)", unionSymbol, member.unionVariantName())
     }
 
+    /**
+     * Renders deserialization logic for a message with the `initial-response` type.
+     */
+    private fun renderDeserializeInitialResponse(ctx: ProtocolGenerator.GenerationContext, outputShape: StructureShape, writer: KotlinWriter) {
+        // A custom function which only deserializes the initial response members
+        val initialResponseDeserializeFn = sdg.payloadDeserializer(ctx, outputShape, outputShape.initialResponseMembers)
+
+        writer.write(
+            "val firstMessage = frames.#T(1).#T()",
+            RuntimeTypes.KotlinxCoroutines.Flow.take,
+            RuntimeTypes.KotlinxCoroutines.Flow.single,
+        )
+        writer.write("val firstMessageType = firstMessage.type()")
+        writer.openBlock(
+            "val events = if (firstMessageType is #T.Event && firstMessageType.shapeType == \"initial-response\") {",
+            RuntimeTypes.AwsEventStream.MessageType,
+        )
+        // Deserialize into `initialResponse`, then apply it to the actual response builder
+        writer.write("val initialResponse = #T(firstMessage.payload)", initialResponseDeserializeFn)
+        writer.withBlock("builder.apply {", "}") {
+            outputShape.initialResponseMembers.forEach { member ->
+                writer.write("#1L = initialResponse.#1L", member.defaultName())
+            }
+        }
+            .write("frames")
+            .closeAndOpenBlock("} else {")
+            .write(
+                "#T(#T(firstMessage), frames)",
+                RuntimeTypes.KotlinxCoroutines.Flow.merge,
+                RuntimeTypes.KotlinxCoroutines.Flow.flowOf,
+            )
+            .closeBlock("}")
+    }
+
+    /**
+     * Get all the shape's members which aren't an event stream
+     */
+    private val StructureShape.initialResponseMembers get() = members().filter {
+        val targetShape = ctx.model.getShape(it.target).getOrNull()
+        targetShape?.hasTrait<StreamingTrait>() == false
+    }
+
     private fun renderDeserializeExplicitEventPayloadMember(
         ctx: ProtocolGenerator.GenerationContext,
         member: MemberShape,

codegen/smithy-aws-kotlin-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/protocols/eventstream/EventStreamParserGenerator.kt

@@ -12,9 +12,11 @@ import software.amazon.smithy.kotlin.codegen.lang.KotlinTypes
 import software.amazon.smithy.kotlin.codegen.model.*
 import software.amazon.smithy.kotlin.codegen.rendering.protocol.ProtocolGenerator
 import software.amazon.smithy.kotlin.codegen.rendering.serde.*
+import software.amazon.smithy.kotlin.codegen.utils.getOrNull
 import software.amazon.smithy.model.shapes.*
 import software.amazon.smithy.model.traits.EventHeaderTrait
 import software.amazon.smithy.model.traits.EventPayloadTrait
+import software.amazon.smithy.model.traits.StreamingTrait
 
 /**
  * Implements rendering serialize implementation for event streams implemented using the
@@ -76,10 +78,35 @@ class EventStreamSerializerGenerator(
         )
 
         val encodeFn = encodeEventStreamMessage(ctx, op, streamShape)
-        writer.withBlock("val messages = stream", "") {
-            write(".#T(::#T)", RuntimeTypes.KotlinxCoroutines.Flow.map, encodeFn)
-            write(".#T(context)", RuntimeTypes.AwsEventStream.sign)
-            write(".#T()", RuntimeTypes.AwsEventStream.encode)
+
+        writer.write("")
+        val initialRequestMembers = input.initialRequestMembers(ctx)
+        if (ctx.protocol.isRpcBoundProtocol && initialRequestMembers.isNotEmpty()) {
+            val serializerFn = sdg.payloadSerializer(ctx, input, initialRequestMembers)
+
+            writer.withBlock("val initialRequest = buildMessage {", "}") {
+                writer.write("addHeader(\":message-type\", HeaderValue.String(\"event\"))")
+                writer.write("addHeader(\":event-type\", HeaderValue.String(\"initial-request\"))")
+                writer.write("payload = #T(input)", serializerFn)
+            }
+
+            writer.withBlock(
+                "val messages = #T(#T(initialRequest), stream.#T(::#T))",
+                "",
+                RuntimeTypes.KotlinxCoroutines.Flow.merge,
+                RuntimeTypes.KotlinxCoroutines.Flow.flowOf,
+                RuntimeTypes.KotlinxCoroutines.Flow.map,
+                encodeFn,
+            ) {
+                write(".#T(context)", RuntimeTypes.AwsEventStream.sign)
+                write(".#T()", RuntimeTypes.AwsEventStream.encode)
+            }
+        } else {
+            writer.withBlock("val messages = stream", "") {
+                write(".#T(::#T)", RuntimeTypes.KotlinxCoroutines.Flow.map, encodeFn)
+                write(".#T(context)", RuntimeTypes.AwsEventStream.sign)
+                write(".#T()", RuntimeTypes.AwsEventStream.encode)
+            }
         }
 
         writer.write("")
@@ -195,4 +222,12 @@ class EventStreamSerializerGenerator(
     private fun KotlinWriter.addStringHeader(name: String, value: String) {
         write("addHeader(#S, #T.String(#S))", name, RuntimeTypes.AwsEventStream.HeaderValue, value)
     }
+
+    /**
+     * Get all the shape's members which aren't an event stream
+     */
+    private fun StructureShape.initialRequestMembers(ctx: ProtocolGenerator.GenerationContext) = members().filter {
+        val targetShape = ctx.model.getShape(it.target).getOrNull()
+        targetShape?.hasTrait<StreamingTrait>() == false
+    }
 }

codegen/smithy-aws-kotlin-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/protocols/eventstream/EventStreamSerializerGenerator.kt

@@ -16,7 +16,6 @@ aws.sdk.kotlin.codegen.customization.BoxServices
 aws.sdk.kotlin.codegen.customization.glacier.GlacierBodyChecksum
 aws.sdk.kotlin.codegen.customization.machinelearning.MachineLearningEndpointCustomization
 aws.sdk.kotlin.codegen.customization.BackfillOptionalAuth
-aws.sdk.kotlin.codegen.customization.RemoveEventStreamOperations
 aws.sdk.kotlin.codegen.customization.flexiblechecksums.FlexibleChecksumsRequest
 aws.sdk.kotlin.codegen.customization.flexiblechecksums.FlexibleChecksumsResponse
 aws.sdk.kotlin.codegen.customization.route53.TrimResourcePrefix

codegen/smithy-aws-kotlin-codegen/src/main/resources/META-INF/services/software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration

@@ -19,24 +19,27 @@ dependencies {
 data class EventStreamTest(
     val projectionName: String,
     val protocolName: String,
+    val modelTemplate: File,
 ) {
     val model: File
         get() = buildDir.resolve("$projectionName/model.smithy")
 }
 
 val tests = listOf(
-    EventStreamTest("restJson1", "restJson1"),
+    EventStreamTest("restJson1", "restJson1", file("event-stream-model-template.smithy")),
+    EventStreamTest("awsJson11", "awsJson1_1", file("event-stream-initial-request-response.smithy")),
 )
 
-fun fillInModel(output: File, protocolName: String) {
-    val template = file("event-stream-model-template.smithy")
+fun fillInModel(output: File, protocolName: String, template: File) {
     val input = template.readText()
     val opTraits = when (protocolName) {
         "restJson1", "restXml" -> """@http(method: "POST", uri: "/test-eventstream", code: 200)"""
         else -> ""
     }
-    val replaced = input.replace("{protocol-name}", protocolName)
+    val replaced = input
+        .replace("{protocol-name}", protocolName)
         .replace("{op-traits}", opTraits)
+
     output.parentFile.mkdirs()
     output.writeText(replaced)
 }
@@ -60,7 +63,7 @@ codegen {
 
             smithyKotlinPlugin {
                 serviceShapeId = testServiceShapeId
-                packageName = "aws.sdk.kotlin.test.eventstream.${test.protocolName.toLowerCase()}"
+                packageName = "aws.sdk.kotlin.test.eventstream.${test.projectionName.toLowerCase()}"
                 packageVersion = "1.0"
                 buildSettings {
                     generateFullProject = false
@@ -77,7 +80,7 @@ codegen {
 
 tasks.named("generateSmithyBuildConfig") {
     doFirst {
-        tests.forEach { test -> fillInModel(test.model, test.protocolName) }
+        tests.forEach { test -> fillInModel(test.model, test.protocolName, test.modelTemplate) }
     }
 }