Enable streaming support for AwsJson (#2207)

* Enable streaming support for AwsJson

Signed-off-by: Daniele Ahmed <[email protected]>

codegen-client-test/build.gradle.kts

@@ -89,6 +89,7 @@ val allCodegenTests = "../codegen-core/common-test-models".let { commonModels ->
         ),
         CodegenTest("aws.protocoltests.json#TestService", "endpoint-rules"),
         CodegenTest("com.aws.example.rust#PokemonService", "pokemon-service-client", imports = listOf("$commonModels/pokemon.smithy", "$commonModels/pokemon-common.smithy")),
+        CodegenTest("com.aws.example.rust#PokemonService", "pokemon-service-awsjson-client", imports = listOf("$commonModels/pokemon-awsjson.smithy", "$commonModels/pokemon-common.smithy")),
     )
 }
 

codegen-core/common-test-models/pokemon-awsjson.smithy

@@ -0,0 +1,104 @@
+$version: "1.0"
+
+// TODO(https://github.com/awslabs/smithy-rs/issues/2215)
+// This is a temporary model to test AwsJson 1.0 with @streaming.
+// This model will be removed when protocol tests support @streaming.
+
+namespace com.aws.example.rust
+
+use aws.protocols#awsJson1_0
+use smithy.framework#ValidationException
+use com.aws.example#PokemonSpecies
+use com.aws.example#GetServerStatistics
+use com.aws.example#DoNothing
+use com.aws.example#CheckHealth
+
+/// The Pokémon Service allows you to retrieve information about Pokémon species.
+@title("Pokémon Service")
+@awsJson1_0
+service PokemonService {
+    version: "2021-12-01",
+    operations: [
+        GetServerStatistics,
+        DoNothing,
+        CapturePokemon,
+        CheckHealth
+    ],
+}
+
+/// Capture Pokémons via event streams.
+@http(uri: "/capture-pokemon-event/{region}", method: "POST")
+operation CapturePokemon {
+    input: CapturePokemonEventsInput,
+    output: CapturePokemonEventsOutput,
+    errors: [UnsupportedRegionError, ThrottlingError, ValidationException]
+}
+
+@input
+structure CapturePokemonEventsInput {
+    @httpPayload
+    events: AttemptCapturingPokemonEvent,
+
+    @httpLabel
+    @required
+    region: String,
+}
+
+@output
+structure CapturePokemonEventsOutput {
+    @httpPayload
+    events: CapturePokemonEvents,
+}
+
+@streaming
+union AttemptCapturingPokemonEvent {
+    event: CapturingEvent,
+    masterball_unsuccessful: MasterBallUnsuccessful,
+}
+
+structure CapturingEvent {
+    @eventPayload
+    payload: CapturingPayload,
+}
+
+structure CapturingPayload {
+    name: String,
+    pokeball: String,
+}
+
+@streaming
+union CapturePokemonEvents {
+    event: CaptureEvent,
+    invalid_pokeball: InvalidPokeballError,
+    throttlingError: ThrottlingError,
+}
+
+structure CaptureEvent {
+    @eventHeader
+    name: String,
+    @eventHeader
+    captured: Boolean,
+    @eventHeader
+    shiny: Boolean,
+    @eventPayload
+    pokedex_update: Blob,
+}
+
+@error("server")
+structure UnsupportedRegionError {
+    @required
+    region: String,
+}
+
+@error("client")
+structure InvalidPokeballError {
+    @required
+    pokeball: String,
+}
+@error("server")
+structure MasterBallUnsuccessful {
+    message: String,
+}
+
+@error("client")
+structure ThrottlingError {}

codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/AwsJson.kt

@@ -25,6 +25,7 @@ import software.amazon.smithy.rust.codegen.core.smithy.protocols.parse.Structure
 import software.amazon.smithy.rust.codegen.core.smithy.protocols.serialize.JsonSerializerGenerator
 import software.amazon.smithy.rust.codegen.core.smithy.protocols.serialize.StructuredDataSerializerGenerator
 import software.amazon.smithy.rust.codegen.core.util.inputShape
+import software.amazon.smithy.rust.codegen.core.util.isStreaming
 
 sealed class AwsJsonVersion {
     abstract val value: String
@@ -50,7 +51,13 @@ class AwsJsonHttpBindingResolver(
 
     private fun bindings(shape: ToShapeId) =
         shape.let { model.expectShape(it.toShapeId()) }.members()
-            .map { HttpBindingDescriptor(it, HttpLocation.DOCUMENT, "document") }
+            .map {
+                if (it.isStreaming(model)) {
+                    HttpBindingDescriptor(it, HttpLocation.PAYLOAD, "document")
+                } else {
+                    HttpBindingDescriptor(it, HttpLocation.DOCUMENT, "document")
+                }
+            }
             .toList()
 
     override fun httpTrait(operationShape: OperationShape): HttpTrait = httpTrait

codegen-server-test/build.gradle.kts

@@ -98,6 +98,11 @@ val allCodegenTests = "../codegen-core/common-test-models".let { commonModels ->
             "pokemon-service-server-sdk",
             imports = listOf("$commonModels/pokemon.smithy", "$commonModels/pokemon-common.smithy"),
         ),
+        CodegenTest(
+            "com.aws.example.rust#PokemonService",
+            "pokemon-service-awsjson-server-sdk",
+            imports = listOf("$commonModels/pokemon-awsjson.smithy", "$commonModels/pokemon-common.smithy"),
+        ),
     )
 }