event streams and operation errors not necessarily disjoint

Signed-off-by: Daniele Ahmed <[email protected]>
smithy-lang · Jul 8, 2022 · 0a150ba · 0a150ba
1 parent 2c4db07
commit 0a150ba
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 13 deletions.
diff --git a/codegen-server-test/model/pokemon.smithy b/codegen-server-test/model/pokemon.smithy
@@ -27,7 +27,7 @@ resource PokemonSpecies {
 operation CapturePokemonOperation {
     input: CapturePokemonOperationEventsInput,
     output: CapturePokemonOperationEventsOutput,
-    errors: [UnsupportedRegionError, MasterBallUnsuccessful]
+    errors: [UnsupportedRegionError, ThrottlingError]
 }
 
 @input
@@ -53,6 +53,7 @@ union AttemptCapturingPokemonEvent {
 }
 
 structure CapturingEvent {
+    @eventPayload
     payload: CapturingPayload,
 }
 
@@ -65,12 +66,17 @@ structure CapturingPayload {
 union CapturePokemonEvents {
     event: CaptureEvent,
     invalid_pokeball: InvalidPokeballError,
+    throttlingError: ThrottlingError,
 }
 
 structure CaptureEvent {
+    @eventHeader
     name: String,
+    @eventHeader
     captured: Boolean,
+    @eventHeader
     shiny: Boolean,
+    @eventPayload
     pokedex_update: Blob,
 }
 
@@ -89,6 +95,8 @@ structure MasterBallUnsuccessful {
     @required
     message: String,
 }
+@error("client")
+structure ThrottlingError {}
 
 /// Retrieve information about a Pokémon species.
 @readonly

diff --git a/...n/kotlin/software/amazon/smithy/rust/codegen/smithy/transformers/EventStreamNormalizer.kt b/...n/kotlin/software/amazon/smithy/rust/codegen/smithy/transformers/EventStreamNormalizer.kt
@@ -21,28 +21,31 @@ import software.amazon.smithy.rust.codegen.smithy.traits.SyntheticEventStreamUni
 import software.amazon.smithy.rust.codegen.util.expectTrait
 import software.amazon.smithy.rust.codegen.util.hasTrait
 import software.amazon.smithy.rust.codegen.util.isEventStream
+import kotlin.streams.toList
 
 /**
  * Generates synthetic unions to replace the modeled unions for Event Stream types.
  * This allows us to strip out all the error union members once up-front, instead of in each
  * place that does codegen with the unions.
  */
 object EventStreamNormalizer {
-    fun transform(model: Model): Model = ModelTransformer.create().mapShapes(model) { shape ->
-        if (shape is OperationShape) {
-            val newErrors = shape.errors
+    fun transform(model: Model): Model {
+        val transformer = ModelTransformer.create()
+        val operations = model.shapes(OperationShape::class.java).toList()
+        val newShapes = operations.flatMap { operation ->
+            operation.errors
                 .map { model.expectShape(it, StructureShape::class.java) }
                 .map { it.toBuilder().addTrait(SyntheticOperationErrorTrait()).build() }
-            val newShape = shape.toBuilder().errors(newErrors.map { it.id }).build()
-            if (newShape.isEventStream(model)) {
-                addStreamErrorsToOperationErrors(model, newShape)
+        }
+        val modelWithSyntheticErrors = model.toBuilder().addShapes(newShapes).build()
+        return transformer.mapShapes(modelWithSyntheticErrors) { shape ->
+            if (shape is OperationShape && shape.isEventStream(modelWithSyntheticErrors)) {
+                addStreamErrorsToOperationErrors(modelWithSyntheticErrors, shape)
+            } else if (shape is UnionShape && shape.isEventStream()) {
+                syntheticEquivalentEventStreamUnion(modelWithSyntheticErrors, shape)
             } else {
-                newShape
+                shape
             }
-        } else if (shape is UnionShape && shape.isEventStream()) {
-            syntheticEquivalentEventStreamUnion(model, shape)
-        } else {
-            shape
         }
     }
 
@@ -107,6 +110,6 @@ fun OperationShape.allErrors(model: Model): List<Shape> {
 
 class SyntheticOperationErrorTrait : AnnotationTrait(ID, Node.objectNode()) {
     companion object {
-        val ID = ShapeId.from("smithy.api.internal#syntheticOperationErrorTrait")
+        val ID: ShapeId = ShapeId.from("smithy.api.internal#syntheticOperationErrorTrait")
     }
 }