Python: Event Streaming (#2482)

* Add `RustType.Application` and `PythonType.Application` variants

* Use `RustType.Application` for event streaming symbols

* Call symbol provider to get `Receiver` type instead of hardcoding

* Add Python specific event streaming symbol provider

* Add event streaming wrapper generator

* Generate correct type information for event streaming types

* Add `CapturePokemon` operation to Python Pokemon service

* Add `InternalServerError` variant to all event streaming errors

* Use `PythonServerCargoDependency` for PyO3Asyncio imports

* Return an attribute error instead of panicking in `IntoPy` impls of wrappers

* Add `Sync` bound to `new` methods of wrappers

* Revert "Add `InternalServerError` variant to all event streaming errors"

This reverts commit b610cb2.

* Add `PythonServerEventStreamErrorGenerator` to generate Python specific error types for unions

* Try to extract error type or inner type from incoming streaming value and ignore the value otherwise for now

* Allow missing type-stubs for `aiohttp`

* Propogate modelled errors through stream

* Raise modelled exceptions rather than sending through stream

* Allow senders from Python side to raise modelled exceptions

* Update `EventStreamSymbolProviderTest` to use `Application` type instead of `Opaque` type

* Fix `ktlint` issues

* Group up common variables in `codegenScope`

* Document `RustType.Application`

* Format `codegen-server-test/python/model/pokemon.smithy`

* Use `tokio-stream` crate instead of `futures` crate

* Use a better error message if user tries to reuse a stream

* Add some details about event streams to example Pokemon service

---------

Co-authored-by: Matteo Bigoi <1781140+crisidev@users.noreply.github.com>

Author: unexge

codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/smithy/EventStreamSymbolProviderTest.kt

@@ -15,6 +15,7 @@ import software.amazon.smithy.rust.codegen.client.testutil.testClientRustSetting
 import software.amazon.smithy.rust.codegen.core.rustlang.RustType
 import software.amazon.smithy.rust.codegen.core.smithy.CodegenTarget
 import software.amazon.smithy.rust.codegen.core.smithy.EventStreamSymbolProvider
+import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
 import software.amazon.smithy.rust.codegen.core.smithy.SymbolVisitor
 import software.amazon.smithy.rust.codegen.core.smithy.rustType
 import software.amazon.smithy.rust.codegen.core.smithy.transformers.OperationNormalizer
@@ -60,8 +61,17 @@ class EventStreamSymbolProviderTest {
         val inputType = provider.toSymbol(inputStream).rustType()
         val outputType = provider.toSymbol(outputStream).rustType()
 
-        inputType shouldBe RustType.Opaque("EventStreamSender<crate::types::SomeStream, crate::types::error::SomeStreamError>", "aws_smithy_http::event_stream")
-        outputType shouldBe RustType.Opaque("Receiver<crate::types::SomeStream, crate::types::error::SomeStreamError>", "aws_smithy_http::event_stream")
+        val someStream = RustType.Opaque("SomeStream", "crate::types")
+        val someStreamError = RustType.Opaque("SomeStreamError", "crate::types::error")
+
+        inputType shouldBe RustType.Application(
+            RuntimeType.eventStreamSender(TestRuntimeConfig).toSymbol().rustType(),
+            listOf(someStream, someStreamError),
+        )
+        outputType shouldBe RustType.Application(
+            RuntimeType.eventStreamReceiver(TestRuntimeConfig).toSymbol().rustType(),
+            listOf(someStream, someStreamError),
+        )
     }
 
     @Test

codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/rustlang/RustType.kt

@@ -172,6 +172,18 @@ sealed class RustType {
     }
 
     data class Opaque(override val name: kotlin.String, override val namespace: kotlin.String? = null) : RustType()
+
+    /**
+     * Represents application of a Rust type with the given arguments.
+     *
+     * For example, we can represent `HashMap<String, i64>` as
+     * `RustType.Application(RustType.Opaque("HashMap"), listOf(RustType.String, RustType.Integer(64)))`.
+     * This helps us to separate the type and the arguments which is useful in methods like [qualifiedName].
+     */
+    data class Application(val type: RustType, val args: List<RustType>) : RustType() {
+        override val name = type.name
+        override val namespace = type.namespace
+    }
 }
 
 /**
@@ -242,7 +254,10 @@ fun RustType.render(fullyQualified: Boolean = true): String {
                 "&${this.lifetime?.let { "'$it" } ?: ""} ${this.member.render(fullyQualified)}"
             }
         }
-
+        is RustType.Application -> {
+            val args = this.args.joinToString(", ") { it.render(fullyQualified) }
+            "${this.name}<$args>"
+        }
         is RustType.Option -> "${this.name}<${this.member.render(fullyQualified)}>"
         is RustType.Box -> "${this.name}<${this.member.render(fullyQualified)}>"
         is RustType.Dyn -> "${this.name} ${this.member.render(fullyQualified)}"

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

@@ -11,7 +11,6 @@ import software.amazon.smithy.model.shapes.OperationShape
 import software.amazon.smithy.model.shapes.Shape
 import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency
 import software.amazon.smithy.rust.codegen.core.rustlang.RustType
-import software.amazon.smithy.rust.codegen.core.rustlang.render
 import software.amazon.smithy.rust.codegen.core.rustlang.stripOuter
 import software.amazon.smithy.rust.codegen.core.smithy.traits.SyntheticInputTrait
 import software.amazon.smithy.rust.codegen.core.smithy.traits.SyntheticOutputTrait
@@ -48,15 +47,15 @@ class EventStreamSymbolProvider(
                 } else {
                     symbolForEventStreamError(unionShape)
                 }
-                val errorFmt = error.rustType().render(fullyQualified = true)
-                val innerFmt = initial.rustType().stripOuter<RustType.Option>().render(fullyQualified = true)
+                val errorT = error.rustType()
+                val innerT = initial.rustType().stripOuter<RustType.Option>()
                 val isSender = (shape.isInputEventStream(model) && target == CodegenTarget.CLIENT) ||
                     (shape.isOutputEventStream(model) && target == CodegenTarget.SERVER)
                 val outer = when (isSender) {
-                    true -> "EventStreamSender<$innerFmt, $errorFmt>"
-                    else -> "Receiver<$innerFmt, $errorFmt>"
+                    true -> RuntimeType.eventStreamSender(runtimeConfig).toSymbol().rustType()
+                    else -> RuntimeType.eventStreamReceiver(runtimeConfig).toSymbol().rustType()
                 }
-                val rustType = RustType.Opaque(outer, "aws_smithy_http::event_stream")
+                val rustType = RustType.Application(outer, listOf(innerT, errorT))
                 return initial.toBuilder()
                     .name(rustType.name)
                     .rustType(rustType)

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

@@ -218,6 +218,9 @@ data class RuntimeType(val path: String, val dependency: RustDependency? = null)
         val Bool = std.resolve("primitive::bool")
         val TryFrom = stdConvert.resolve("TryFrom")
         val Vec = std.resolve("vec::Vec")
+        val Arc = std.resolve("sync::Arc")
+        val Send = std.resolve("marker::Send")
+        val Sync = std.resolve("marker::Sync")
 
         // external cargo dependency types
         val Bytes = CargoDependency.Bytes.toType().resolve("Bytes")
@@ -277,6 +280,7 @@ data class RuntimeType(val path: String, val dependency: RustDependency? = null)
         fun document(runtimeConfig: RuntimeConfig): RuntimeType = smithyTypes(runtimeConfig).resolve("Document")
         fun retryErrorKind(runtimeConfig: RuntimeConfig) = smithyTypes(runtimeConfig).resolve("retry::ErrorKind")
         fun eventStreamReceiver(runtimeConfig: RuntimeConfig): RuntimeType = smithyHttp(runtimeConfig).resolve("event_stream::Receiver")
+        fun eventStreamSender(runtimeConfig: RuntimeConfig): RuntimeType = smithyHttp(runtimeConfig).resolve("event_stream::EventStreamSender")
         fun errorMetadata(runtimeConfig: RuntimeConfig) = smithyTypes(runtimeConfig).resolve("error::ErrorMetadata")
         fun errorMetadataBuilder(runtimeConfig: RuntimeConfig) = smithyTypes(runtimeConfig).resolve("error::metadata::Builder")
         fun provideErrorMetadataTrait(runtimeConfig: RuntimeConfig) = smithyTypes(runtimeConfig).resolve("error::metadata::ProvideErrorMetadata")

codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/http/HttpBindingGenerator.kt

@@ -30,6 +30,7 @@ import software.amazon.smithy.rust.codegen.core.rustlang.RustType
 import software.amazon.smithy.rust.codegen.core.rustlang.RustWriter
 import software.amazon.smithy.rust.codegen.core.rustlang.Writable
 import software.amazon.smithy.rust.codegen.core.rustlang.asOptional
+import software.amazon.smithy.rust.codegen.core.rustlang.qualifiedName
 import software.amazon.smithy.rust.codegen.core.rustlang.render
 import software.amazon.smithy.rust.codegen.core.rustlang.rust
 import software.amazon.smithy.rust.codegen.core.rustlang.rustBlock
@@ -222,7 +223,7 @@ class HttpBindingGenerator(
                     // Streaming unions are Event Streams and should be handled separately
                     val target = model.expectShape(binding.member.target)
                     if (target is UnionShape) {
-                        bindEventStreamOutput(operationShape, target)
+                        bindEventStreamOutput(operationShape, outputT, target)
                     } else {
                         deserializeStreamingBody(binding)
                     }
@@ -243,22 +244,22 @@ class HttpBindingGenerator(
         }
     }
 
-    private fun RustWriter.bindEventStreamOutput(operationShape: OperationShape, targetShape: UnionShape) {
+    private fun RustWriter.bindEventStreamOutput(operationShape: OperationShape, outputT: Symbol, targetShape: UnionShape) {
         val unmarshallerConstructorFn = EventStreamUnmarshallerGenerator(
             protocol,
             codegenContext,
             operationShape,
             targetShape,
         ).render()
+        val receiver = outputT.rustType().qualifiedName()
         rustTemplate(
             """
             let unmarshaller = #{unmarshallerConstructorFn}();
             let body = std::mem::replace(body, #{SdkBody}::taken());
-            Ok(#{Receiver}::new(unmarshaller, body))
+            Ok($receiver::new(unmarshaller, body))
             """,
             "SdkBody" to RuntimeType.sdkBody(runtimeConfig),
             "unmarshallerConstructorFn" to unmarshallerConstructorFn,
-            "Receiver" to RuntimeType.eventStreamReceiver(runtimeConfig),
         )
     }
 

codegen-server-test/python/model/pokemon.smithy

@@ -1,19 +1,18 @@
 /// TODO(https://github.com/awslabs/smithy-rs/issues/1508)
-/// $econcile this model with the main one living inside codegen-server-test/model/pokemon.smithy
+/// Reconcile this model with the main one living inside codegen-server-test/model/pokemon.smithy
 /// once the Python implementation supports Streaming and Union shapes.
 $version: "1.0"
 
 namespace com.aws.example.python
 
 use aws.protocols#restJson1
+use com.aws.example#CheckHealth
+use com.aws.example#DoNothing
+use com.aws.example#GetServerStatistics
 use com.aws.example#PokemonSpecies
 use com.aws.example#Storage
-use com.aws.example#GetServerStatistics
-use com.aws.example#DoNothing
-use com.aws.example#CheckHealth
 use smithy.framework#ValidationException
 
-
 /// The Pokémon Service allows you to retrieve information about Pokémon species.
 @title("Pokémon Service")
 @restJson1
@@ -23,11 +22,93 @@ service PokemonService {
     operations: [
         GetServerStatistics
         DoNothing
+        CapturePokemon
         CheckHealth
         StreamPokemonRadio
-    ],
+    ]
+}
+
+/// 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 {}
+
 /// Fetch the radio song from the database and stream it back as a playable audio.
 @readonly
 @http(uri: "/radio", method: "GET")

codegen-server/python/src/main/kotlin/software/amazon/smithy/rust/codegen/server/python/smithy/PythonEventStreamSymbolProvider.kt

@@ -0,0 +1,104 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.rust.codegen.server.python.smithy
+
+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.rust.codegen.core.rustlang.CargoDependency
+import software.amazon.smithy.rust.codegen.core.rustlang.RustType
+import software.amazon.smithy.rust.codegen.core.rustlang.stripOuter
+import software.amazon.smithy.rust.codegen.core.smithy.RuntimeConfig
+import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
+import software.amazon.smithy.rust.codegen.core.smithy.RustSymbolProvider
+import software.amazon.smithy.rust.codegen.core.smithy.WrappingSymbolProvider
+import software.amazon.smithy.rust.codegen.core.smithy.rustType
+import software.amazon.smithy.rust.codegen.core.smithy.traits.SyntheticInputTrait
+import software.amazon.smithy.rust.codegen.core.smithy.traits.SyntheticOutputTrait
+import software.amazon.smithy.rust.codegen.core.smithy.transformers.eventStreamErrors
+import software.amazon.smithy.rust.codegen.core.util.getTrait
+import software.amazon.smithy.rust.codegen.core.util.isEventStream
+import software.amazon.smithy.rust.codegen.core.util.isOutputEventStream
+import software.amazon.smithy.rust.codegen.core.util.toPascalCase
+
+/**
+ * Symbol provider for Python that maps event streaming member shapes to their respective Python wrapper types.
+ *
+ * For example given a model:
+ * ```smithy
+ * @input
+ * structure CapturePokemonInput {
+ *     @httpPayload
+ *     events: AttemptCapturingPokemonEvent,
+ * }
+ *
+ * @streaming
+ * union AttemptCapturingPokemonEvent {
+ *     ...
+ * }
+ * ```
+ * for the member shape `CapturePokemonInput$events` it will return a symbol that points to
+ * `crate::python_event_stream::CapturePokemonInputEventsReceiver`.
+ */
+class PythonEventStreamSymbolProvider(
+    private val runtimeConfig: RuntimeConfig,
+    base: RustSymbolProvider,
+) : WrappingSymbolProvider(base) {
+    override fun toSymbol(shape: Shape): Symbol {
+        val initial = super.toSymbol(shape)
+
+        // We only want to wrap with Event Stream types when dealing with member shapes
+        if (shape !is MemberShape || !shape.isEventStream(model)) {
+            return initial
+        }
+
+        // We can only wrap the type if it's either an input or an output that used in an operation
+        model.expectShape(shape.container).let { maybeInputOutput ->
+            val operationId = maybeInputOutput.getTrait<SyntheticInputTrait>()?.operation
+                ?: maybeInputOutput.getTrait<SyntheticOutputTrait>()?.operation
+            operationId?.let { model.expectShape(it, OperationShape::class.java) }
+        } ?: return initial
+
+        val unionShape = model.expectShape(shape.target).asUnionShape().get()
+        val error = if (unionShape.eventStreamErrors().isEmpty()) {
+            RuntimeType.smithyHttp(runtimeConfig).resolve("event_stream::MessageStreamError").toSymbol()
+        } else {
+            symbolForEventStreamError(unionShape)
+        }
+        val inner = initial.rustType().stripOuter<RustType.Option>()
+        val innerSymbol = Symbol.builder().name(inner.name).rustType(inner).build()
+        val containerName = shape.container.name
+        val memberName = shape.memberName.toPascalCase()
+        val outer = when (shape.isOutputEventStream(model)) {
+            true -> "${containerName}${memberName}EventStreamSender"
+            else -> "${containerName}${memberName}Receiver"
+        }
+        val rustType = RustType.Opaque(outer, PythonServerRustModule.PythonEventStream.fullyQualifiedPath())
+        return Symbol.builder()
+            .name(rustType.name)
+            .rustType(rustType)
+            .addReference(innerSymbol)
+            .addReference(error)
+            .addDependency(CargoDependency.smithyHttp(runtimeConfig).withFeature("event-stream"))
+            .addDependency(PythonServerCargoDependency.TokioStream)
+            .addDependency(PythonServerCargoDependency.PyO3Asyncio.withFeature("unstable-streams"))
+            .build()
+    }
+
+    companion object {
+        data class EventStreamSymbol(val innerT: RustType, val errorT: RustType)
+
+        fun parseSymbol(symbol: Symbol): EventStreamSymbol {
+            check(symbol.references.size >= 2) {
+                "`PythonEventStreamSymbolProvider` adds inner type and error type as references to resulting symbol"
+            }
+            val innerT = symbol.references[0].symbol.rustType()
+            val errorT = symbol.references[1].symbol.rustType()
+            return EventStreamSymbol(innerT, errorT)
+        }
+    }
+}

codegen-server/python/src/main/kotlin/software/amazon/smithy/rust/codegen/server/python/smithy/PythonServerCargoDependency.kt

@@ -18,6 +18,7 @@ object PythonServerCargoDependency {
     val PyO3: CargoDependency = CargoDependency("pyo3", CratesIo("0.17"))
     val PyO3Asyncio: CargoDependency = CargoDependency("pyo3-asyncio", CratesIo("0.17"), features = setOf("attributes", "tokio-runtime"))
     val Tokio: CargoDependency = CargoDependency("tokio", CratesIo("1.20.1"), features = setOf("full"))
+    val TokioStream: CargoDependency = CargoDependency("tokio-stream", CratesIo("0.1.12"))
     val Tracing: CargoDependency = CargoDependency("tracing", CratesIo("0.1"))
     val Tower: CargoDependency = CargoDependency("tower", CratesIo("0.4"))
     val TowerHttp: CargoDependency = CargoDependency("tower-http", CratesIo("0.3"), features = setOf("trace"))