Add server codegen flag addValidationExceptionToConstrainedOperations (…

…#3803)

The Server SDK requires the model to include
`aws.smithy.framework#ValidationException` in each operation that can
access a constrained member shape. This becomes problematic when
generating the server SDK for a model not owned by the team creating the
SDK, as they cannot easily modify the model.

This PR introduces a codegen flag,
`addValidationExceptionToConstrainedOperations`. When set in
`smithy-build-template.json`, this flag will automatically add
`ValidationException` to operations that require it but do not already
list it among their errors.

Closes Issue:
[3802](#3802)

Sample `smithy-build-template.json`
```
"plugins": {
            "rust-server-codegen": {
                "service": "ServiceToGenerateSDKFor",
                "module": "amzn-sample-server-sdk",
                "codegen": {
                    "addValidationExceptionToConstrainedOperations": true,
                }
            }
        }
```

---------

Co-authored-by: Fahad Zubair <fahadzub@amazon.com>

.changelog/4619777.md

@@ -0,0 +1,26 @@
+---
+applies_to: ["server"]
+authors: ["drganjoo"]
+references: ["smithy-rs#3803"]
+breaking: false
+new_feature: true
+bug_fix: false
+---
+Setting the `addValidationExceptionToConstrainedOperations` codegen flag adds `aws.smithy.framework#ValidationException` to operations with constrained inputs that do not already have this exception added.
+
+Sample `smithy-build-template.json`:
+
+```
+{
+    "...",
+    "plugins": {
+        "rust-server-codegen": {
+            "service": "ServiceToGenerateSDKFor",
+                "module": "amzn-sample-server-sdk",
+                "codegen": {
+                    "addValidationExceptionToConstrainedOperations": true,
+                }
+        }
+    }
+}
+```

codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/testutil/CodegenIntegrationTest.kt

@@ -29,6 +29,109 @@ data class IntegrationTestParams(
     val cargoCommand: String? = null,
 )
 
+/**
+ * A helper class to allow setting `codegen` object keys to be passed to the `additionalSettings`
+ * field of `IntegrationTestParams`.
+ *
+ * Usage:
+ *
+ * ```kotlin
+ * serverIntegrationTest(
+ *     model,
+ *     IntegrationTestParams(
+ *         additionalSettings =
+ *             ServerAdditionalSettings.builder()
+ *                 .generateCodegenComments()
+ *                 .publicConstrainedTypes()
+ *                 .toObjectNode()
+ * )),
+ * ```
+ */
+sealed class AdditionalSettings {
+    abstract fun toObjectNode(): ObjectNode
+
+    abstract class CoreAdditionalSettings protected constructor(val settings: List<AdditionalSettings>) : AdditionalSettings() {
+        override fun toObjectNode(): ObjectNode {
+            val merged =
+                settings.map { it.toObjectNode() }
+                    .reduce { acc, next -> acc.merge(next) }
+
+            return ObjectNode.builder()
+                .withMember("codegen", merged)
+                .build()
+        }
+
+        abstract class Builder<T : CoreAdditionalSettings> : AdditionalSettings() {
+            protected val settings = mutableListOf<AdditionalSettings>()
+
+            fun generateCodegenComments(debugMode: Boolean = true): Builder<T> {
+                settings.add(GenerateCodegenComments(debugMode))
+                return this
+            }
+
+            abstract fun build(): T
+
+            override fun toObjectNode(): ObjectNode = build().toObjectNode()
+        }
+
+        // Core settings that are common to both Servers and Clients should be defined here.
+        data class GenerateCodegenComments(val debugMode: Boolean) : AdditionalSettings() {
+            override fun toObjectNode(): ObjectNode =
+                ObjectNode.builder()
+                    .withMember("debugMode", debugMode)
+                    .build()
+        }
+    }
+}
+
+class ClientAdditionalSettings private constructor(settings: List<AdditionalSettings>) :
+    AdditionalSettings.CoreAdditionalSettings(settings) {
+        class Builder : CoreAdditionalSettings.Builder<ClientAdditionalSettings>() {
+            override fun build(): ClientAdditionalSettings = ClientAdditionalSettings(settings)
+        }
+
+        // Additional settings that are specific to client generation should be defined here.
+
+        companion object {
+            fun builder() = Builder()
+        }
+    }
+
+class ServerAdditionalSettings private constructor(settings: List<AdditionalSettings>) :
+    AdditionalSettings.CoreAdditionalSettings(settings) {
+        class Builder : CoreAdditionalSettings.Builder<ServerAdditionalSettings>() {
+            fun publicConstrainedTypes(enabled: Boolean = true): Builder {
+                settings.add(PublicConstrainedTypes(enabled))
+                return this
+            }
+
+            fun addValidationExceptionToConstrainedOperations(enabled: Boolean = true): Builder {
+                settings.add(AddValidationExceptionToConstrainedOperations(enabled))
+                return this
+            }
+
+            override fun build(): ServerAdditionalSettings = ServerAdditionalSettings(settings)
+        }
+
+        private data class PublicConstrainedTypes(val enabled: Boolean) : AdditionalSettings() {
+            override fun toObjectNode(): ObjectNode =
+                ObjectNode.builder()
+                    .withMember("publicConstrainedTypes", enabled)
+                    .build()
+        }
+
+        private data class AddValidationExceptionToConstrainedOperations(val enabled: Boolean) : AdditionalSettings() {
+            override fun toObjectNode(): ObjectNode =
+                ObjectNode.builder()
+                    .withMember("addValidationExceptionToConstrainedOperations", enabled)
+                    .build()
+        }
+
+        companion object {
+            fun builder() = Builder()
+        }
+    }
+
 /**
  * Run cargo test on a true, end-to-end, codegen product of a given model.
  */

codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/ServerCodegenVisitor.kt

@@ -86,7 +86,7 @@ import software.amazon.smithy.rust.codegen.server.smithy.generators.protocol.Ser
 import software.amazon.smithy.rust.codegen.server.smithy.generators.protocol.ServerProtocolTestGenerator
 import software.amazon.smithy.rust.codegen.server.smithy.protocols.ServerProtocolLoader
 import software.amazon.smithy.rust.codegen.server.smithy.traits.isReachableFromOperationInput
-import software.amazon.smithy.rust.codegen.server.smithy.transformers.AttachValidationExceptionToConstrainedOperationInputsInAllowList
+import software.amazon.smithy.rust.codegen.server.smithy.transformers.AttachValidationExceptionToConstrainedOperationInputs
 import software.amazon.smithy.rust.codegen.server.smithy.transformers.ConstrainedMemberTransform
 import software.amazon.smithy.rust.codegen.server.smithy.transformers.RecursiveConstraintViolationBoxer
 import software.amazon.smithy.rust.codegen.server.smithy.transformers.RemoveEbsModelValidationException
@@ -204,8 +204,8 @@ open class ServerCodegenVisitor(
             // Remove the EBS model's own `ValidationException`, which collides with `smithy.framework#ValidationException`
             .let(RemoveEbsModelValidationException::transform)
             // Attach the `smithy.framework#ValidationException` error to operations whose inputs are constrained,
-            // if they belong to a service in an allowlist
-            .let(AttachValidationExceptionToConstrainedOperationInputsInAllowList::transform)
+            // if either the operation belongs to a service in the allowlist, or the codegen flag to add the exception has been set.
+            .let { AttachValidationExceptionToConstrainedOperationInputs.transform(it, settings) }
             // Tag aggregate shapes reachable from operation input
             .let(ShapesReachableFromOperationInputTagger::transform)
             // Normalize event stream operations

codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/ServerRustSettings.kt

@@ -96,13 +96,15 @@ data class ServerCodegenConfig(
      *  able to define the converters in their Rust application code.
      */
     val experimentalCustomValidationExceptionWithReasonPleaseDoNotUse: String? = defaultExperimentalCustomValidationExceptionWithReasonPleaseDoNotUse,
+    val addValidationExceptionToConstrainedOperations: Boolean = DEFAULT_ADD_VALIDATION_EXCEPTION_TO_CONSTRAINED_OPERATIONS,
 ) : CoreCodegenConfig(
         formatTimeoutSeconds, debugMode,
     ) {
     companion object {
         private const val DEFAULT_PUBLIC_CONSTRAINED_TYPES = true
         private const val DEFAULT_IGNORE_UNSUPPORTED_CONSTRAINTS = false
         private val defaultExperimentalCustomValidationExceptionWithReasonPleaseDoNotUse = null
+        private const val DEFAULT_ADD_VALIDATION_EXCEPTION_TO_CONSTRAINED_OPERATIONS = false
 
         fun fromCodegenConfigAndNode(
             coreCodegenConfig: CoreCodegenConfig,
@@ -114,6 +116,7 @@ data class ServerCodegenConfig(
                 publicConstrainedTypes = node.get().getBooleanMemberOrDefault("publicConstrainedTypes", DEFAULT_PUBLIC_CONSTRAINED_TYPES),
                 ignoreUnsupportedConstraints = node.get().getBooleanMemberOrDefault("ignoreUnsupportedConstraints", DEFAULT_IGNORE_UNSUPPORTED_CONSTRAINTS),
                 experimentalCustomValidationExceptionWithReasonPleaseDoNotUse = node.get().getStringMemberOrDefault("experimentalCustomValidationExceptionWithReasonPleaseDoNotUse", defaultExperimentalCustomValidationExceptionWithReasonPleaseDoNotUse),
+                addValidationExceptionToConstrainedOperations = node.get().getBooleanMemberOrDefault("addValidationExceptionToConstrainedOperations", DEFAULT_ADD_VALIDATION_EXCEPTION_TO_CONSTRAINED_OPERATIONS),
             )
         } else {
             ServerCodegenConfig(

codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/transformers/AttachValidationExceptionToConstrainedOperationInputsInAllowList.kt → codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/transformers/AttachValidationExceptionToConstrainedOperationInputs.kt

@@ -8,14 +8,41 @@ package software.amazon.smithy.rust.codegen.server.smithy.transformers
 import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.shapes.EnumShape
 import software.amazon.smithy.model.shapes.OperationShape
+import software.amazon.smithy.model.shapes.ServiceShape
 import software.amazon.smithy.model.shapes.SetShape
 import software.amazon.smithy.model.shapes.ShapeId
 import software.amazon.smithy.model.transform.ModelTransformer
 import software.amazon.smithy.rust.codegen.core.smithy.DirectedWalker
 import software.amazon.smithy.rust.codegen.core.util.inputShape
+import software.amazon.smithy.rust.codegen.server.smithy.ServerRustSettings
 import software.amazon.smithy.rust.codegen.server.smithy.customizations.SmithyValidationExceptionConversionGenerator
 import software.amazon.smithy.rust.codegen.server.smithy.hasConstraintTrait
 
+private fun addValidationExceptionToMatchingServiceShapes(
+    model: Model,
+    filterPredicate: (ServiceShape) -> Boolean,
+): Model {
+    val walker = DirectedWalker(model)
+    val operationsWithConstrainedInputWithoutValidationException =
+        model.serviceShapes
+            .filter(filterPredicate)
+            .flatMap { it.operations }
+            .map { model.expectShape(it, OperationShape::class.java) }
+            .filter { operationShape ->
+                walker.walkShapes(operationShape.inputShape(model))
+                    .any { it is SetShape || it is EnumShape || it.hasConstraintTrait() }
+            }
+            .filter { !it.errors.contains(SmithyValidationExceptionConversionGenerator.SHAPE_ID) }
+
+    return ModelTransformer.create().mapShapes(model) { shape ->
+        if (shape is OperationShape && operationsWithConstrainedInputWithoutValidationException.contains(shape)) {
+            shape.toBuilder().addError(SmithyValidationExceptionConversionGenerator.SHAPE_ID).build()
+        } else {
+            shape
+        }
+    }
+}
+
 /**
  * Attach the `smithy.framework#ValidationException` error to operations whose inputs are constrained, if they belong
  * to a service in an allowlist.
@@ -32,7 +59,7 @@ import software.amazon.smithy.rust.codegen.server.smithy.hasConstraintTrait
  *  `disableDefaultValidation` set to `true`, allowing service owners to map from constraint violations to operation errors.
  */
 object AttachValidationExceptionToConstrainedOperationInputsInAllowList {
-    private val sherviceShapeIdAllowList =
+    private val serviceShapeIdAllowList =
         setOf(
             // These we currently generate server SDKs for.
             ShapeId.from("aws.protocoltests.restjson#RestJson"),
@@ -50,25 +77,47 @@ object AttachValidationExceptionToConstrainedOperationInputsInAllowList {
         )
 
     fun transform(model: Model): Model {
-        val walker = DirectedWalker(model)
-        val operationsWithConstrainedInputWithoutValidationException =
-            model.serviceShapes
-                .filter { sherviceShapeIdAllowList.contains(it.toShapeId()) }
-                .flatMap { it.operations }
-                .map { model.expectShape(it, OperationShape::class.java) }
-                .filter { operationShape ->
-                    // Walk the shapes reachable via this operation input.
-                    walker.walkShapes(operationShape.inputShape(model))
-                        .any { it is SetShape || it is EnumShape || it.hasConstraintTrait() }
-                }
-                .filter { !it.errors.contains(SmithyValidationExceptionConversionGenerator.SHAPE_ID) }
+        return addValidationExceptionToMatchingServiceShapes(
+            model,
+        ) { serviceShapeIdAllowList.contains(it.toShapeId()) }
+    }
+}
 
-        return ModelTransformer.create().mapShapes(model) { shape ->
-            if (shape is OperationShape && operationsWithConstrainedInputWithoutValidationException.contains(shape)) {
-                shape.toBuilder().addError(SmithyValidationExceptionConversionGenerator.SHAPE_ID).build()
-            } else {
-                shape
-            }
+/**
+ * Attach the `smithy.framework#ValidationException` error to operations with constrained inputs if the
+ * codegen flag `addValidationExceptionToConstrainedOperations` has been set.
+ */
+object AttachValidationExceptionToConstrainedOperationInputsBasedOnCodegenFlag {
+    fun transform(
+        model: Model,
+        settings: ServerRustSettings,
+    ): Model {
+        if (!settings.codegenConfig.addValidationExceptionToConstrainedOperations) {
+            return model
         }
+
+        val service = settings.getService(model)
+
+        return addValidationExceptionToMatchingServiceShapes(
+            model,
+        ) { it == service }
+    }
+}
+
+/**
+ * Attaches the `smithy.framework#ValidationException` error to operations with constrained inputs
+ * if either of the following conditions is met:
+ * 1. The operation belongs to a service in the allowlist.
+ * 2. The codegen flag `addValidationExceptionToConstrainedOperations` has been set.
+ *
+ * The error is only attached if the operation does not already have `ValidationException` added.
+ */
+object AttachValidationExceptionToConstrainedOperationInputs {
+    fun transform(
+        model: Model,
+        settings: ServerRustSettings,
+    ): Model {
+        val allowListTransformedModel = AttachValidationExceptionToConstrainedOperationInputsInAllowList.transform(model)
+        return AttachValidationExceptionToConstrainedOperationInputsBasedOnCodegenFlag.transform(allowListTransformedModel, settings)
     }
 }

codegen-server/src/test/kotlin/software/amazon/smithy/rust/codegen/server/smithy/customizations/AddValidationExceptionToConstrainedOperationsTest.kt

@@ -0,0 +1,85 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.rust.codegen.server.smithy.customizations
+
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.assertThrows
+import software.amazon.smithy.codegen.core.CodegenException
+import software.amazon.smithy.rust.codegen.core.testutil.IntegrationTestParams
+import software.amazon.smithy.rust.codegen.core.testutil.ServerAdditionalSettings
+import software.amazon.smithy.rust.codegen.core.testutil.asSmithyModel
+import software.amazon.smithy.rust.codegen.server.smithy.testutil.serverIntegrationTest
+
+/**
+ * Tests whether the server `codegen` flag `addValidationExceptionToConstrainedOperations` works as expected.
+ */
+internal class AddValidationExceptionToConstrainedOperationsTest {
+    private val testModelWithValidationExceptionImported =
+        """
+        namespace test
+
+        use smithy.framework#ValidationException
+        use aws.protocols#restJson1
+        use aws.api#data
+
+        @restJson1
+        service ConstrainedService {
+            operations: [SampleOperationWithValidation, SampleOperationWithoutValidation]
+        }
+
+        @http(uri: "/anOperationWithValidation", method: "POST")
+        operation SampleOperationWithValidation {
+            output: SampleInputOutput
+            input: SampleInputOutput
+            errors: [ValidationException, ErrorWithMemberConstraint]
+        }
+        @http(uri: "/anOperationWithoutValidation", method: "POST")
+        operation SampleOperationWithoutValidation {
+            output: SampleInputOutput
+            input: SampleInputOutput
+            errors: []
+        }
+        structure SampleInputOutput {
+            constrainedInteger : RangedInteger
+            @range(min: 2, max:100)
+            constrainedMemberInteger : RangedInteger
+            patternString : PatternString
+        }
+        @pattern("^[a-m]+${'$'}")
+        string PatternString
+        @range(min: 0, max:1000)
+        integer RangedInteger
+
+        @error("server")
+        structure ErrorWithMemberConstraint {
+            @range(min: 100, max: 999)
+            statusCode: Integer
+        }
+        """.asSmithyModel(smithyVersion = "2")
+
+    @Test
+    fun `without setting the codegen flag, the model should fail to compile`() {
+        assertThrows<CodegenException> {
+            serverIntegrationTest(
+                testModelWithValidationExceptionImported,
+                IntegrationTestParams(),
+            )
+        }
+    }
+
+    @Test
+    fun `operations that do not have ValidationException will automatically have one added to them`() {
+        serverIntegrationTest(
+            testModelWithValidationExceptionImported,
+            IntegrationTestParams(
+                additionalSettings =
+                    ServerAdditionalSettings.builder()
+                        .addValidationExceptionToConstrainedOperations()
+                        .toObjectNode(),
+            ),
+        )
+    }
+}