Server streaming body

Signed-off-by: Guy Margalit <[email protected]>
smithy-lang · Feb 1, 2022 · bc571b3 · bc571b3
1 parent 384c492
commit bc571b3
Show file tree

Hide file tree

Showing 7 changed files with 138 additions and 95 deletions.
diff --git a/...re/amazon/smithy/rust/codegen/server/smithy/generators/ServerOperationHandlerGenerator.kt b/...re/amazon/smithy/rust/codegen/server/smithy/generators/ServerOperationHandlerGenerator.kt
@@ -18,6 +18,8 @@ import software.amazon.smithy.rust.codegen.server.smithy.protocols.ServerHttpPro
 import software.amazon.smithy.rust.codegen.smithy.CodegenContext
 import software.amazon.smithy.rust.codegen.smithy.RuntimeType
 import software.amazon.smithy.rust.codegen.smithy.generators.error.errorSymbol
+import software.amazon.smithy.rust.codegen.util.hasStreamingMember
+import software.amazon.smithy.rust.codegen.util.inputShape
 import software.amazon.smithy.rust.codegen.util.outputShape
 
 /**
@@ -132,13 +134,18 @@ class ServerOperationHandlerGenerator(
         } else {
             symbolProvider.toSymbol(operation.outputShape(model)).fullName
         }
+        val streamingBodyTraitBounds = if (operation.inputShape(model).hasStreamingMember(model)) {
+            "\n B: Into<#{SmithyHttpServer}::ByteStream>,"
+        } else {
+            ""
+        }
         return """
             $inputFn
             Fut: std::future::Future<Output = $outputType> + Send,
-            B: $serverCrate::HttpBody + Send + 'static,
+            B: $serverCrate::HttpBody + Send + 'static, $streamingBodyTraitBounds
             B::Data: Send,
             B::Error: Into<$serverCrate::BoxError>,
             $serverCrate::rejection::SmithyRejection: From<<B as $serverCrate::HttpBody>::Error>
-        """
+        """.trimIndent()
     }
 }
diff --git a/...azon/smithy/rust/codegen/server/smithy/generators/protocol/ServerProtocolTestGenerator.kt b/...azon/smithy/rust/codegen/server/smithy/generators/protocol/ServerProtocolTestGenerator.kt
@@ -212,16 +212,16 @@ class ServerProtocolTestGenerator(
 
         rustTemplate(
             """
-                ##[allow(unused_mut)] let mut http_request = http::Request::builder()
-                    .uri("${httpRequestTestCase.uri}")
-                """,
-                *codegenScope
+            ##[allow(unused_mut)] let mut http_request = http::Request::builder()
+                .uri("${httpRequestTestCase.uri}")
+            """,
+            *codegenScope
         )
         for (header in httpRequestTestCase.headers) {
             rust(".header(${header.key.dq()}, ${header.value.dq()})")
         }
         rustTemplate(
-    """
+            """
             .body(#{SmithyHttpServer}::Body::from(#{Bytes}::from_static(b${httpRequestTestCase.body.orNull()?.dq()})))
             .unwrap();
             """,
@@ -363,19 +363,19 @@ class ServerProtocolTestGenerator(
         basicCheck(
             requireHeaders,
             rustWriter,
-           "required_headers",
-           actualExpression,
-           "require_headers"
+            "required_headers",
+            actualExpression,
+            "require_headers"
         )
     }
 
     private fun checkForbidHeaders(rustWriter: RustWriter, actualExpression: String, forbidHeaders: List<String>) {
         basicCheck(
             forbidHeaders,
             rustWriter,
-           "forbidden_headers",
+            "forbidden_headers",
             actualExpression,
-           "forbid_headers"
+            "forbid_headers"
         )
     }
 
@@ -525,14 +525,14 @@ class ServerProtocolTestGenerator(
             FailingTest(RestJson, "RestJsonSupportsNegativeInfinityFloatInputs", Action.Response),
             FailingTest(RestJson, "RestJsonStreamingTraitsWithBlob", Action.Request),
             FailingTest(RestJson, "RestJsonStreamingTraitsWithNoBlobBody", Action.Request),
-            FailingTest(RestJson, "RestJsonStreamingTraitsWithBlob", Action.Response),
-            FailingTest(RestJson, "RestJsonStreamingTraitsWithNoBlobBody", Action.Response),
+            // FailingTest(RestJson, "RestJsonStreamingTraitsWithBlob", Action.Response),
+            // FailingTest(RestJson, "RestJsonStreamingTraitsWithNoBlobBody", Action.Response),
             FailingTest(RestJson, "RestJsonStreamingTraitsRequireLengthWithBlob", Action.Request),
             FailingTest(RestJson, "RestJsonStreamingTraitsRequireLengthWithNoBlobBody", Action.Request),
             FailingTest(RestJson, "RestJsonStreamingTraitsRequireLengthWithBlob", Action.Response),
-            FailingTest(RestJson, "RestJsonStreamingTraitsRequireLengthWithNoBlobBody", Action.Response),
+            // FailingTest(RestJson, "RestJsonStreamingTraitsRequireLengthWithNoBlobBody", Action.Response),
             FailingTest(RestJson, "RestJsonStreamingTraitsWithMediaTypeWithBlob", Action.Request),
-            FailingTest(RestJson, "RestJsonStreamingTraitsWithMediaTypeWithBlob", Action.Response),
+            // FailingTest(RestJson, "RestJsonStreamingTraitsWithMediaTypeWithBlob", Action.Response),
             FailingTest(RestJson, "RestJsonHttpWithEmptyBlobPayload", Action.Request),
             FailingTest(RestJson, "RestJsonHttpWithEmptyStructurePayload", Action.Request),
 
@@ -603,8 +603,9 @@ class ServerProtocolTestGenerator(
                 ).asObjectNode().get()
             ).build()
         private fun fixRestJsonAllQueryStringTypes(testCase: HttpRequestTestCase): HttpRequestTestCase =
-             testCase.toBuilder().params(
-                 Node.parse("""{
+            testCase.toBuilder().params(
+                Node.parse(
+                    """{
                     "queryString": "Hello there",
                     "queryStringList": ["a", "b", "c"],
                     "queryStringSet": ["a", "b", "c"],
@@ -643,16 +644,19 @@ class ServerProtocolTestGenerator(
                         "Enum": ["Foo"],
                         "EnumList": ["Foo", "Baz", "Bar"]
                     }
-                }""".trimMargin()).asObjectNode().get()
-             ).build()
+                }""".trimMargin()
+                ).asObjectNode().get()
+            ).build()
         private fun fixRestJsonQueryStringEscaping(testCase: HttpRequestTestCase): HttpRequestTestCase =
             testCase.toBuilder().params(
-                Node.parse("""{
-                   "queryString": "%:/?#[]@!${'$'}&'()*+,;=😹",
-                   "queryParamsMapOfStringList": {
-                       "String": ["%:/?#[]@!${'$'}&'()*+,;=😹"]
-                   }
-               }""".trimMargin()).asObjectNode().get()
+                Node.parse(
+                    """{
+                    "queryString": "%:/?#[]@!${'$'}&'()*+,;=😹",
+                    "queryParamsMapOfStringList": {
+                        "String": ["%:/?#[]@!${'$'}&'()*+,;=😹"]
+                    }
+               }""".trimMargin()
+                ).asObjectNode().get()
             ).build()
         // This test assumes that errors in responses are identified by an `X-Amzn-Errortype` header with the error shape name.
         // However, Smithy specifications for AWS protocols that serialize to JSON recommend that new server implementations