Autoapply the input trait for any services that are missing it

smithy-lang · Oct 17, 2023 · 61194a5 · 61194a5
1 parent 92b0c9f
commit 61194a5
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/...ware/amazon/smithy/rust/codegen/client/smithy/customizations/IdempotencyTokenGenerator.kt b/...ware/amazon/smithy/rust/codegen/client/smithy/customizations/IdempotencyTokenGenerator.kt
@@ -26,7 +26,7 @@ import software.amazon.smithy.rust.codegen.core.util.inputShape
 
 class IdempotencyTokenGenerator(
     codegenContext: CodegenContext,
-    operationShape: OperationShape,
+    private val operationShape: OperationShape,
 ) : OperationCustomization() {
     private val model = codegenContext.model
     private val runtimeConfig = codegenContext.runtimeConfig
@@ -56,7 +56,7 @@ class IdempotencyTokenGenerator(
             is OperationSection.AdditionalRuntimePlugins -> writable {
                 section.addOperationRuntimePlugin(this) {
                     if (!symbolProvider.toSymbol(idempotencyTokenMember).isOptional()) {
-                        UNREACHABLE("top level input members are always optional")
+                        UNREACHABLE("top level input members are always optional. $operationShape")
                     }
                     // An idempotency token is optional. If the user didn't specify a token
                     // then we'll generate one and set it.

diff --git a/...otlin/software/amazon/smithy/rust/codegen/core/smithy/transformers/OperationNormalizer.kt b/...otlin/software/amazon/smithy/rust/codegen/core/smithy/transformers/OperationNormalizer.kt
@@ -10,9 +10,11 @@ import software.amazon.smithy.model.shapes.OperationShape
 import software.amazon.smithy.model.shapes.Shape
 import software.amazon.smithy.model.shapes.ShapeId
 import software.amazon.smithy.model.shapes.StructureShape
+import software.amazon.smithy.model.traits.InputTrait
 import software.amazon.smithy.model.transform.ModelTransformer
 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.util.hasTrait
 import software.amazon.smithy.rust.codegen.core.util.orNull
 import software.amazon.smithy.rust.codegen.core.util.rename
 import java.util.Optional
@@ -106,6 +108,11 @@ object OperationNormalizer {
         val inputShapeBuilder = operation.input.map { shapeId ->
             model.expectShape(shapeId, StructureShape::class.java).toBuilder().rename(inputId)
         }.orElse(empty(inputId))
+        // There are still shapes missing the input trait. If we don't apply this, we'll get bad results from the
+        // nullability index
+        if (!inputShapeBuilder.build().hasTrait<InputTrait>()) {
+            inputShapeBuilder.addTrait(InputTrait())
+        }
         return inputShapeBuilder.addTrait(
             SyntheticInputTrait(
                 operation = operation.id,