Skip to content

Commit

Permalink
Support server event streams
Browse files Browse the repository at this point in the history
* Server event streams
* Rename EventStreamInput to EventStreamSender
* Make event stream errors optional
* Pokemon service model updated
* Pokemon server event handler
* Pokemon client to test event streams
* EventStreamDecorator to make optional using SigV4 signing

Closes: #1157

Signed-off-by: Daniele Ahmed <ahmeddan@amazon.de>
  • Loading branch information
82marbag committed Jun 21, 2022
1 parent a3c7902 commit c4f1809
Show file tree
Hide file tree
Showing 27 changed files with 544 additions and 98 deletions.
6 changes: 6 additions & 0 deletions CHANGELOG.next.toml
Original file line number Diff line number Diff line change
Expand Up @@ -47,3 +47,9 @@ the prefix will be stripped automatically.
references = ["aws-sdk-rust#554"]
meta = { "breaking" = false, "tada" = false, "bug" = true }
author = "Velfi"

[[smithy-rs]]
message = "Rename EventStreamInput to EventStreamSender"
references = ["smithy-rs#1157"]
meta = { "breaking" = true, "tada" = false, "bug" = false }
author = "82marbag"
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ import software.amazon.smithy.rust.codegen.smithy.customizations.RetryConfigDeco
import software.amazon.smithy.rust.codegen.smithy.customizations.SleepImplDecorator
import software.amazon.smithy.rust.codegen.smithy.customizations.TimeoutConfigDecorator
import software.amazon.smithy.rust.codegen.smithy.customize.CombinedCodegenDecorator
import software.amazon.smithy.rust.codegen.smithy.customize.EventStreamDecorator
import software.amazon.smithy.rustsdk.customize.apigateway.ApiGatewayDecorator
import software.amazon.smithy.rustsdk.customize.auth.DisabledAuthDecorator
import software.amazon.smithy.rustsdk.customize.ec2.Ec2Decorator
Expand All @@ -24,7 +25,7 @@ val DECORATORS = listOf(
RegionDecorator(),
AwsEndpointDecorator(),
UserAgentDecorator(),
SigV4SigningDecorator(),
EventStreamDecorator(listOf(SigV4SigningDecorator())),
RetryPolicyDecorator(),
IntegrationTestDecorator(),
AwsFluentClientDecorator(),
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,9 +22,9 @@ import software.amazon.smithy.rust.codegen.rustlang.writable
import software.amazon.smithy.rust.codegen.smithy.CodegenContext
import software.amazon.smithy.rust.codegen.smithy.RuntimeConfig
import software.amazon.smithy.rust.codegen.smithy.RuntimeType
import software.amazon.smithy.rust.codegen.smithy.customize.EventStreamDecorator
import software.amazon.smithy.rust.codegen.smithy.customize.OperationCustomization
import software.amazon.smithy.rust.codegen.smithy.customize.OperationSection
import software.amazon.smithy.rust.codegen.smithy.customize.RustCodegenDecorator
import software.amazon.smithy.rust.codegen.smithy.generators.config.ConfigCustomization
import software.amazon.smithy.rust.codegen.smithy.generators.config.ServiceConfig
import software.amazon.smithy.rust.codegen.smithy.letIf
Expand All @@ -42,11 +42,11 @@ import software.amazon.smithy.rust.codegen.util.isInputEventStream
* - sets a default `OperationSigningConfig` A future enhancement will customize this for specific services that need
* different behavior.
*/
class SigV4SigningDecorator : RustCodegenDecorator {
class SigV4SigningDecorator : EventStreamDecorator(listOf()) {
override val name: String = "SigV4Signing"
override val order: Byte = 0

private fun applies(codegenContext: CodegenContext): Boolean = codegenContext.serviceShape.hasTrait<SigV4Trait>()
override fun applies(codegenContext: CodegenContext): Boolean = codegenContext.serviceShape.hasTrait<SigV4Trait>()

override fun configCustomizations(
codegenContext: CodegenContext,
Expand Down
40 changes: 39 additions & 1 deletion codegen-server-test/model/pokemon.smithy
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ use aws.protocols#restJson1
service PokemonService {
version: "2021-12-01",
resources: [PokemonSpecies],
operations: [GetServerStatistics, EmptyOperation],
operations: [GetServerStatistics, EmptyOperation, CapturePokemonOperation],
}

/// A Pokémon species forms the basis for at least one Pokémon.
Expand All @@ -22,6 +22,44 @@ resource PokemonSpecies {
read: GetPokemonSpecies,
}

/// Capture Pokémons via event streams
@http(uri: "/capture-pokemon-event", method: "POST")
operation CapturePokemonOperation {
input: CapturePokemonOperationEventsInput,
output: CapturePokemonOperationEventsOutput,
}

@input
structure CapturePokemonOperationEventsInput {
@httpPayload
events: AttemptCapturingPokemonEvent,
}

@output
structure CapturePokemonOperationEventsOutput {
@httpPayload
events: CapturePokemonEvents,
}

@streaming
union AttemptCapturingPokemonEvent {
event: CapturingEvent,
}

structure CapturingEvent {
name: String,
pokeball: String,
}

@streaming
union CapturePokemonEvents {
event: CaptureEvent,
}

structure CaptureEvent {
name: String,
}

/// Retrieve information about a Pokémon species.
@readonly
@http(uri: "/pokemon-species/{name}", method: "GET")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ class PythonCodegenServerPlugin : SmithyBuildPlugin {
override fun execute(context: PluginContext) {
// Suppress extremely noisy logs about reserved words
Logger.getLogger(ReservedWordSymbolProvider::class.java.name).level = Level.OFF
// Discover [RustCodegenDecorators] on the classpath. [RustCodegenDectorator] return different types of
// Discover [RustCodegenDecorators] on the classpath. [RustCodegenDecorator] return different types of
// customization. A customization is a function of:
// - location (e.g. the mutate section of an operation)
// - context (e.g. the of the operation)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ import software.amazon.smithy.rust.codegen.smithy.StreamingShapeSymbolProvider
import software.amazon.smithy.rust.codegen.smithy.SymbolVisitor
import software.amazon.smithy.rust.codegen.smithy.SymbolVisitorConfig
import software.amazon.smithy.rust.codegen.smithy.customize.CombinedCodegenDecorator
import software.amazon.smithy.rust.codegen.smithy.generators.CodegenTarget
import java.util.logging.Level
import java.util.logging.Logger

Expand Down Expand Up @@ -62,7 +63,7 @@ class RustCodegenServerPlugin : SmithyBuildPlugin {
SymbolVisitor(model, serviceShape = serviceShape, config = symbolVisitorConfig)
// Generate different types for EventStream shapes (e.g. transcribe streaming)
.let {
EventStreamSymbolProvider(symbolVisitorConfig.runtimeConfig, it, model)
EventStreamSymbolProvider(symbolVisitorConfig.runtimeConfig, it, model, CodegenTarget.SERVER)
}
// Generate [ByteStream] instead of `Blob` for streaming binary shapes (e.g. S3 GetObject)
.let { StreamingShapeSymbolProvider(it, model) }
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -505,13 +505,10 @@ private class ServerHttpBoundProtocolTraitImplGenerator(
// Fallback to the default code of `http::response::Builder`, 200.

operationShape.outputShape(model).findStreamingMember(model)?.let {
val memberName = symbolProvider.toMemberName(it)
rustTemplate(
"""
let body = #{SmithyHttpServer}::body::to_boxed(#{SmithyHttpServer}::body::Body::wrap_stream(output.$memberName));
""",
*codegenScope,
)
val payloadGenerator = HttpBoundProtocolPayloadGenerator(codegenContext, protocol, httpMessageType = HttpMessageType.RESPONSE)
withBlockTemplate("let body = #{SmithyHttpServer}::body::boxed(#{SmithyHttpServer}::body::Body::wrap_stream(", "));", *codegenScope) {
payloadGenerator.generatePayload(this, "output", operationShape)
}
} ?: run {
val payloadGenerator = HttpBoundProtocolPayloadGenerator(codegenContext, protocol, httpMessageType = HttpMessageType.RESPONSE)
withBlockTemplate("let body = #{SmithyHttpServer}::body::to_boxed(", ");", *codegenScope) {
Expand Down Expand Up @@ -682,29 +679,29 @@ private class ServerHttpBoundProtocolTraitImplGenerator(
HttpLocation.HEADER -> writable { serverRenderHeaderParser(this, binding, operationShape) }
HttpLocation.PREFIX_HEADERS -> writable { serverRenderPrefixHeadersParser(this, binding, operationShape) }
HttpLocation.PAYLOAD -> {
return if (binding.member.isStreaming(model)) {
writable {
val structureShapeHandler: RustWriter.(String) -> Unit = { body ->
rust("#T($body)", structuredDataParser.payloadParser(binding.member))
}
val errorSymbol = getDeserializePayloadErrorSymbol(binding)
val deserializer = httpBindingGenerator.generateDeserializePayloadFn(
binding,
errorSymbol,
structuredHandler = structureShapeHandler
)
return writable {
if (binding.member.isStreaming(model)) {
rustTemplate(
"""
{
let body = request.take_body().ok_or(#{RequestRejection}::BodyAlreadyExtracted)?;
Some(body.into())
let bytes = #{Hyper}::body::to_bytes(body).await?;
Some(#{Deserializer}(&mut bytes.into())?)
}
""".trimIndent(),
""",
"Deserializer" to deserializer,
*codegenScope
)
}
} else {
val structureShapeHandler: RustWriter.(String) -> Unit = { body ->
rust("#T($body)", structuredDataParser.payloadParser(binding.member))
}
val errorSymbol = getDeserializePayloadErrorSymbol(binding)
val deserializer = httpBindingGenerator.generateDeserializePayloadFn(
binding,
errorSymbol,
structuredHandler = structureShapeHandler
)
writable {
} else {
rustTemplate(
"""
{
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,21 +14,25 @@ import software.amazon.smithy.rust.codegen.rustlang.CargoDependency
import software.amazon.smithy.rust.codegen.rustlang.RustType
import software.amazon.smithy.rust.codegen.rustlang.render
import software.amazon.smithy.rust.codegen.rustlang.stripOuter
import software.amazon.smithy.rust.codegen.smithy.generators.CodegenTarget
import software.amazon.smithy.rust.codegen.smithy.generators.error.errorSymbol
import software.amazon.smithy.rust.codegen.smithy.traits.SyntheticInputTrait
import software.amazon.smithy.rust.codegen.smithy.traits.SyntheticOutputTrait
import software.amazon.smithy.rust.codegen.util.getTrait
import software.amazon.smithy.rust.codegen.util.isEventStream
import software.amazon.smithy.rust.codegen.util.isInputEventStream
import software.amazon.smithy.rust.codegen.util.isOutputEventStream

/**
* Wrapping symbol provider to wrap modeled types with the aws-smithy-http Event Stream send/receive types.
*/
class EventStreamSymbolProvider(
private val runtimeConfig: RuntimeConfig,
base: RustSymbolProvider,
private val model: Model
private val model: Model,
private val target: CodegenTarget,
) : WrappingSymbolProvider(base) {
private val smithyEventStream = CargoDependency.SmithyEventStream(runtimeConfig)
override fun toSymbol(shape: Shape): Symbol {
val initial = super.toSymbol(shape)

Expand All @@ -42,21 +46,30 @@ class EventStreamSymbolProvider(
}
// If we find an operation shape, then we can wrap the type
if (operationShape != null) {
val error = operationShape.errorSymbol(this).toSymbol()
val error = if (operationShape.errors.isNotEmpty()) {
operationShape.errorSymbol(this).toSymbol()
} else {
RuntimeType("MessageStreamError", smithyEventStream, "aws_smithy_http::event_stream")
.toSymbol()
}
val errorFmt = error.rustType().render(fullyQualified = true)
val innerFmt = initial.rustType().stripOuter<RustType.Option>().render(fullyQualified = true)
val outer = when (shape.isInputEventStream(model)) {
true -> "EventStreamInput<$innerFmt>"
val isSender = (shape.isInputEventStream(model) && target == CodegenTarget.CLIENT) ||
(shape.isOutputEventStream(model) && target == CodegenTarget.SERVER)
val outer = when (isSender) {
true -> "EventStreamSender<$innerFmt>"
else -> "Receiver<$innerFmt, $errorFmt>"
}
val rustType = RustType.Opaque(outer, "aws_smithy_http::event_stream")
return initial.toBuilder()
val symbol = initial.toBuilder()
.name(rustType.name)
.rustType(rustType)
.addReference(error)
.addReference(initial)
.addDependency(CargoDependency.SmithyHttp(runtimeConfig).withFeature("event-stream"))
.build()
if (operationShape.errors.isNotEmpty()) {
symbol.addReference(error)
}
return symbol.build()
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ import software.amazon.smithy.rust.codegen.rustlang.Attribute.Companion.NonExhau
import software.amazon.smithy.rust.codegen.rustlang.RustReservedWordSymbolProvider
import software.amazon.smithy.rust.codegen.smithy.customizations.ClientCustomizations
import software.amazon.smithy.rust.codegen.smithy.customize.CombinedCodegenDecorator
import software.amazon.smithy.rust.codegen.smithy.generators.CodegenTarget
import java.util.logging.Level
import java.util.logging.Logger

Expand Down Expand Up @@ -49,7 +50,7 @@ class RustCodegenPlugin : SmithyBuildPlugin {
fun baseSymbolProvider(model: Model, serviceShape: ServiceShape, symbolVisitorConfig: SymbolVisitorConfig = DefaultConfig) =
SymbolVisitor(model, serviceShape = serviceShape, config = symbolVisitorConfig)
// Generate different types for EventStream shapes (e.g. transcribe streaming)
.let { EventStreamSymbolProvider(symbolVisitorConfig.runtimeConfig, it, model) }
.let { EventStreamSymbolProvider(symbolVisitorConfig.runtimeConfig, it, model, CodegenTarget.CLIENT) }
// Generate `ByteStream` instead of `Blob` for streaming binary shapes (e.g. S3 GetObject)
.let { StreamingShapeSymbolProvider(it, model) }
// Add Rust attributes (like `#[derive(PartialEq)]`) to generated shapes
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,93 @@
/*
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
* SPDX-License-Identifier: Apache-2.0
*/

package software.amazon.smithy.rust.codegen.smithy.customize

import software.amazon.smithy.model.shapes.OperationShape
import software.amazon.smithy.rust.codegen.rustlang.CargoDependency
import software.amazon.smithy.rust.codegen.rustlang.Writable
import software.amazon.smithy.rust.codegen.rustlang.rustTemplate
import software.amazon.smithy.rust.codegen.rustlang.writable
import software.amazon.smithy.rust.codegen.smithy.CodegenContext
import software.amazon.smithy.rust.codegen.smithy.RuntimeConfig
import software.amazon.smithy.rust.codegen.smithy.RuntimeType
import software.amazon.smithy.rust.codegen.smithy.generators.config.ConfigCustomization
import software.amazon.smithy.rust.codegen.smithy.generators.config.ServiceConfig
import software.amazon.smithy.rust.codegen.util.hasEventStreamOperations

/**
* The EventStreamDecorator:
* - adds a `new_event_stream_signer()` method to `config` to create an Event Stream NoOp signer
* - can be customized by subclassing, see SigV4SigningDecorator
*/
open class EventStreamDecorator(
private val decorators: List<EventStreamDecorator>
) : RustCodegenDecorator {
override val name: String = "EventStreamDecorator"
override val order: Byte = 0

open fun applies(codegenContext: CodegenContext): Boolean = true

override fun configCustomizations(
codegenContext: CodegenContext,
baseCustomizations: List<ConfigCustomization>
): List<ConfigCustomization> {
decorators.forEach {
if (it.applies(codegenContext)) return it.configCustomizations(codegenContext, baseCustomizations)
}
return baseCustomizations + EventStreamSignConfig(
codegenContext.runtimeConfig,
codegenContext.serviceShape.hasEventStreamOperations(codegenContext.model),
)
}

override fun operationCustomizations(
codegenContext: CodegenContext,
operation: OperationShape,
baseCustomizations: List<OperationCustomization>
): List<OperationCustomization> {
decorators.forEach {
if (it.applies(codegenContext)) return it.operationCustomizations(codegenContext, operation, baseCustomizations)
}
return baseCustomizations
}
}

class EventStreamSignConfig(
runtimeConfig: RuntimeConfig,
private val serviceHasEventStream: Boolean,
) : ConfigCustomization() {
private val smithyEventStream = CargoDependency.SmithyEventStream(runtimeConfig)
private val codegenScope = arrayOf(
"NoOpSigner" to RuntimeType("NoOpSigner", smithyEventStream, "aws_smithy_eventstream::frame"),
"SharedPropertyBag" to RuntimeType(
"SharedPropertyBag",
CargoDependency.SmithyHttp(runtimeConfig),
"aws_smithy_http::property_bag"
)
)

override fun section(section: ServiceConfig): Writable {
return when (section) {
is ServiceConfig.ConfigImpl -> writable {
if (serviceHasEventStream) {
rustTemplate(
"""
/// Creates a new Event Stream `SignMessage` implementor.
pub fn new_event_stream_signer(
&self,
_properties: #{SharedPropertyBag}
) -> #{NoOpSigner} {
#{NoOpSigner}{}
}
""",
*codegenScope
)
}
}
else -> emptySection
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -155,7 +155,7 @@ open class CombinedCodegenDecorator(decorators: List<RustCodegenDecorator>) : Ru
.onEach {
logger.info("Adding Codegen Decorator: ${it.javaClass.name}")
}.toList()
return CombinedCodegenDecorator(decorators + RequiredCustomizations() + FluentClientDecorator() + extras)
return CombinedCodegenDecorator(decorators + RequiredCustomizations() + EventStreamDecorator(listOf()) + FluentClientDecorator() + extras)
}
}
}
Loading

0 comments on commit c4f1809

Please sign in to comment.