-
Notifications
You must be signed in to change notification settings - Fork 180
/
EventStreamSymbolProvider.kt
77 lines (72 loc) · 3.98 KB
/
EventStreamSymbolProvider.kt
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
/*
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
* SPDX-License-Identifier: Apache-2.0
*/
package software.amazon.smithy.rust.codegen.smithy
import software.amazon.smithy.codegen.core.Symbol
import software.amazon.smithy.model.Model
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.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.eventStreamErrorSymbol
import software.amazon.smithy.rust.codegen.smithy.traits.SyntheticInputTrait
import software.amazon.smithy.rust.codegen.smithy.traits.SyntheticOutputTrait
import software.amazon.smithy.rust.codegen.smithy.transformers.eventStreamErrors
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 target: CodegenTarget,
) : WrappingSymbolProvider(base) {
private val smithyEventStream = CargoDependency.SmithyEventStream(runtimeConfig)
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)) {
// Determine if the member has a container that is a synthetic input or output
val operationShape = model.expectShape(shape.container).let { maybeInputOutput ->
val operationId = maybeInputOutput.getTrait<SyntheticInputTrait>()?.operation
?: maybeInputOutput.getTrait<SyntheticOutputTrait>()?.operation
operationId?.let { model.expectShape(it, OperationShape::class.java) }
}
// If we find an operation shape, then we can wrap the type
if (operationShape != null) {
val unionShape = model.expectShape(shape.target).asUnionShape().get()
val error = if (target == CodegenTarget.SERVER && unionShape.eventStreamErrors().isEmpty()) {
RuntimeType("MessageStreamError", smithyEventStream, "aws_smithy_http::event_stream").toSymbol()
} else {
unionShape.eventStreamErrorSymbol(model, this, target).toSymbol()
}
val errorFmt = error.rustType().render(fullyQualified = true)
val innerFmt = initial.rustType().stripOuter<RustType.Option>().render(fullyQualified = true)
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>"
}
val rustType = RustType.Opaque(outer, "aws_smithy_http::event_stream")
return initial.toBuilder()
.name(rustType.name)
.rustType(rustType)
.addReference(initial)
.addDependency(CargoDependency.SmithyHttp(runtimeConfig).withFeature("event-stream"))
.addReference(error)
.build()
}
}
return initial
}
}