Unverified Commit 1a7a495d authored by Burak's avatar Burak Committed by GitHub
Browse files

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 b610cb27c7532102e3a3ec15a59da6c56b60c318.

* 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: default avatarMatteo Bigoi <1781140+crisidev@users.noreply.github.com>
parent 33b24bb7

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

+12 −2
Original line number Diff line number Diff line
@@ -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

+16 −1
Original line number Diff line number Diff line
@@ -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

+5 −6
Original line number Diff line number Diff line
@@ -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

+4 −0
Original line number Diff line number Diff line
@@ -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

+5 −4
Original line number Diff line number Diff line
@@ -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),

        )

    }