Reexport EventReceiver in a generated client crate

CHANGELOG.next.toml

@@ -9,4 +9,16 @@
 # message = "Fix typos in module documentation for generated crates"
 # references = ["smithy-rs#920"]
 # meta = { "breaking" = false, "tada" = false, "bug" = false, "target" = "client | server | all"}
-# author = "rcoh"
+# author = "rcoh"
+
+[[smithy-rs]]
+message = "`crate::event_receiver::EventReceiver` is now re-exported as `crate::primitives::event_stream::EventReceiver` when a service supports event stream operations."
+references = ["smithy-rs#3305"]
+meta = { "breaking" = false, "tada" = false, "bug" = true, "target" = "client" }
+author = "ysaito1001"
+
+[[aws-sdk-rust]]
+message = "`crate::event_receiver::EventReceiver` is now re-exported as `crate::primitives::event_stream::EventReceiver` when a service supports event stream operations."
+references = ["smithy-rs#3305"]
+meta = { "breaking" = false, "tada" = false, "bug" = true }
+author = "ysaito1001"

aws/sdk/integration-tests/transcribestreaming/tests/test.rs

@@ -7,6 +7,8 @@ use async_stream::stream;
 use aws_sdk_transcribestreaming::config::{Credentials, Region};
 use aws_sdk_transcribestreaming::error::SdkError;
 use aws_sdk_transcribestreaming::operation::start_stream_transcription::StartStreamTranscriptionOutput;
+#[allow(unused)] // making sure `EventReceiver` is re-exported
+use aws_sdk_transcribestreaming::primitives::event_stream::EventReceiver;
 use aws_sdk_transcribestreaming::primitives::event_stream::{HeaderValue, Message};
 use aws_sdk_transcribestreaming::primitives::Blob;
 use aws_sdk_transcribestreaming::types::error::{AudioStreamError, TranscriptResultStreamError};

codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/customizations/SmithyTypesPubUseExtra.kt

@@ -14,6 +14,7 @@ import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
 import software.amazon.smithy.rust.codegen.core.rustlang.writable
 import software.amazon.smithy.rust.codegen.core.smithy.CodegenContext
 import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
+import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType.Companion.eventReceiver
 import software.amazon.smithy.rust.codegen.core.smithy.RustCrate
 import software.amazon.smithy.rust.codegen.core.util.hasEventStreamMember
 import software.amazon.smithy.rust.codegen.core.util.hasEventStreamOperations
@@ -88,11 +89,13 @@ fun pubUseSmithyPrimitivesEventStream(codegenContext: CodegenContext, model: Mod
     if (codegenContext.serviceShape.hasEventStreamOperations(model)) {
         rustTemplate(
             """
+            pub use #{EventReceiver};
             pub use #{Header};
             pub use #{HeaderValue};
             pub use #{Message};
             pub use #{StrBytes};
             """,
+            "EventReceiver" to eventReceiver(rc),
             "Header" to RuntimeType.smithyTypes(rc).resolve("event_stream::Header"),
             "HeaderValue" to RuntimeType.smithyTypes(rc).resolve("event_stream::HeaderValue"),
             "Message" to RuntimeType.smithyTypes(rc).resolve("event_stream::Message"),

codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/customizations/SmithyTypesPubUseExtraTest.kt

@@ -10,7 +10,6 @@ import software.amazon.smithy.model.Model
 import software.amazon.smithy.rust.codegen.core.rustlang.RustWriter
 import software.amazon.smithy.rust.codegen.core.smithy.CodegenContext
 import software.amazon.smithy.rust.codegen.core.smithy.RustCrate
-import software.amazon.smithy.rust.codegen.core.smithy.generators.StructureGeneratorTest.Companion.model
 import software.amazon.smithy.rust.codegen.core.testutil.asSmithyModel
 import software.amazon.smithy.rust.codegen.core.testutil.generatePluginContext
 import software.amazon.smithy.rust.codegen.core.testutil.testCodegenContext
@@ -25,6 +24,11 @@ class SmithyTypesPubUseExtraTest {
         return """
             namespace test
 
+            service TestService {
+                version: "123"
+                operations: [SomeOperation]
+            }
+
             $additionalShape
             structure SomeStruct {
             }
@@ -46,20 +50,21 @@ class SmithyTypesPubUseExtraTest {
         """.asSmithyModel()
     }
 
-    private val rustCrate: RustCrate
-    private val codegenContext: CodegenContext = testCodegenContext(model)
+    private fun initialize(model: Model): Pair<CodegenContext, RustCrate> {
+        val codegenContext = testCodegenContext(model)
 
-    init {
         val (context, _) = generatePluginContext(
             model,
             runtimeConfig = codegenContext.runtimeConfig,
         )
-        rustCrate = RustCrate(
+        val rustCrate = RustCrate(
             context.fileManifest,
             codegenContext.symbolProvider,
             codegenContext.settings.codegenConfig,
             codegenContext.expectModuleDocProvider(),
         )
+
+        return Pair(codegenContext, rustCrate)
     }
 
     private fun reexportsWithEmptyModel() = reexportsWithMember()
@@ -69,11 +74,12 @@ class SmithyTypesPubUseExtraTest {
         unionMember: String = "",
         additionalShape: String = "",
     ) = RustWriter.root().let { writer ->
-        pubUseSmithyPrimitives(
-            codegenContext,
-            modelWithMember(inputMember, outputMember, unionMember, additionalShape),
-            rustCrate,
-        )(writer)
+        val model = modelWithMember(inputMember, outputMember, unionMember, additionalShape)
+        val props = initialize(model)
+        val context = props.first
+        val rustCrate = props.second
+        pubUseSmithyPrimitives(context, model, rustCrate)(writer)
+        pubUseSmithyPrimitivesEventStream(context, model)(writer)
         writer.toString()
     }
 
@@ -151,4 +157,26 @@ class SmithyTypesPubUseExtraTest {
             streamingTypes,
         )
     }
+
+    @Test
+    fun `it re-exports when a model has event stream`() {
+        val eventStreamTypes =
+            listOf(
+                "crate::event_receiver::EventReceiver",
+                "::aws_smithy_types::event_stream::Header",
+                "::aws_smithy_types::event_stream::HeaderValue",
+                "::aws_smithy_types::event_stream::Message",
+                "::aws_smithy_types::str_bytes::StrBytes",
+            )
+        val eventStreamShape = "@streaming union EventStream { foo: SomeStruct }"
+
+        assertHasReexports(
+            reexportsWithMember(additionalShape = eventStreamShape, inputMember = "m: EventStream"),
+            eventStreamTypes,
+        )
+        assertHasReexports(
+            reexportsWithMember(additionalShape = eventStreamShape, outputMember = "m: EventStream"),
+            eventStreamTypes,
+        )
+    }
 }