Add injection hook for structures from a ServiceShape (#3147)

Analogous to #3001, for `ServiceShape`s ---- _By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice._ Signed-off-by: Daniele Ahmed <[email protected]>
smithy-lang · Nov 8, 2023 · 79a0e72 · 79a0e72
1 parent cc2b947
commit 79a0e72
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 4 deletions.
diff --git a/...src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/ServerCodegenVisitor.kt b/...src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/ServerCodegenVisitor.kt
@@ -624,6 +624,9 @@ open class ServerCodegenVisitor(
 
             ScopeMacroGenerator(codegenContext).render(this)
         }
+
+        codegenDecorator.postprocessServiceGenerateAdditionalStructures(shape)
+            .forEach { structureShape -> this.structureShape(structureShape) }
     }
 
     /**
@@ -649,7 +652,7 @@ open class ServerCodegenVisitor(
             protocolGenerator.renderOperation(this, shape)
         }
 
-        codegenDecorator.postprocessGenerateAdditionalStructures(shape)
+        codegenDecorator.postprocessOperationGenerateAdditionalStructures(shape)
             .forEach { structureShape -> this.structureShape(structureShape) }
     }
 

diff --git a/...lin/software/amazon/smithy/rust/codegen/server/smithy/customize/ServerCodegenDecorator.kt b/...lin/software/amazon/smithy/rust/codegen/server/smithy/customize/ServerCodegenDecorator.kt
@@ -7,6 +7,7 @@ package software.amazon.smithy.rust.codegen.server.smithy.customize
 
 import software.amazon.smithy.build.PluginContext
 import software.amazon.smithy.model.shapes.OperationShape
+import software.amazon.smithy.model.shapes.ServiceShape
 import software.amazon.smithy.model.shapes.ShapeId
 import software.amazon.smithy.model.shapes.StructureShape
 import software.amazon.smithy.rust.codegen.core.smithy.customize.CombinedCoreCodegenDecorator
@@ -41,7 +42,15 @@ interface ServerCodegenDecorator : CoreCodegenDecorator<ServerCodegenContext, Se
      * making the resulting crate not compile (since it will contain more than one struct with the same name).
      * Therefore, ensure that all the structure shapes returned by this method are not in the service's closure.
      */
-    fun postprocessGenerateAdditionalStructures(operationShape: OperationShape): List<StructureShape> = emptyList()
+    fun postprocessOperationGenerateAdditionalStructures(operationShape: OperationShape): List<StructureShape> = emptyList()
+
+    /**
+     * For each service, this hook allows decorators to return a collection of structure shapes that will additionally be generated.
+     * If a structure shape is in the service's closure, note that returning it here will cause for it to be generated more than once,
+     * making the resulting crate not compile (since it will contain more than one struct with the same name).
+     * Therefore, ensure that all the structure shapes returned by this method are not in the service's closure.
+     */
+    fun postprocessServiceGenerateAdditionalStructures(serviceShape: ServiceShape): List<StructureShape> = emptyList()
 
     /**
      * Configuration methods that should be injected into the `${serviceName}Config` struct to allow users to configure
@@ -81,8 +90,11 @@ class CombinedServerCodegenDecorator(decorators: List<ServerCodegenDecorator>) :
             decorator.postprocessValidationExceptionNotAttachedErrorMessage(accumulated)
         }
 
-    override fun postprocessGenerateAdditionalStructures(operationShape: OperationShape): List<StructureShape> =
-        orderedDecorators.flatMap { it.postprocessGenerateAdditionalStructures(operationShape) }
+    override fun postprocessOperationGenerateAdditionalStructures(operationShape: OperationShape): List<StructureShape> =
+        orderedDecorators.flatMap { it.postprocessOperationGenerateAdditionalStructures(operationShape) }
+
+    override fun postprocessServiceGenerateAdditionalStructures(serviceShape: ServiceShape): List<StructureShape> =
+        orderedDecorators.flatMap { it.postprocessServiceGenerateAdditionalStructures(serviceShape) }
 
     override fun configMethods(codegenContext: ServerCodegenContext): List<ConfigMethod> =
         orderedDecorators.flatMap { it.configMethods(codegenContext) }