diff --git a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/CodegenDirector.java b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/CodegenDirector.java index 3b36479138d..795fc54fb00 100644 --- a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/CodegenDirector.java +++ b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/CodegenDirector.java @@ -302,6 +302,19 @@ public void changeStringEnumsToEnumShapes(boolean synthesizeEnumNames) { }); } + /** + * Flattens service-level pagination information into operation pagination traits. + * + * @see ModelTransformer#flattenPaginationInfoIntoOperations(Model, ServiceShape) + */ + public void flattenPaginationInfoIntoOperations() { + transforms.add((model, transformer) -> { + LOGGER.finest("Flattening pagination info into operation traits for directed codegen"); + return transformer.flattenPaginationInfoIntoOperations(model, + model.expectShape(service, ServiceShape.class)); + }); + } + /** * Sets the shapes order for code generation. * diff --git a/smithy-model/src/main/java/software/amazon/smithy/model/transform/FlattenPaginationInfo.java b/smithy-model/src/main/java/software/amazon/smithy/model/transform/FlattenPaginationInfo.java new file mode 100644 index 00000000000..c7e1aa1632d --- /dev/null +++ b/smithy-model/src/main/java/software/amazon/smithy/model/transform/FlattenPaginationInfo.java @@ -0,0 +1,52 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.model.transform; + +import java.util.HashSet; +import java.util.Optional; +import java.util.Set; +import software.amazon.smithy.model.Model; +import software.amazon.smithy.model.knowledge.PaginatedIndex; +import software.amazon.smithy.model.knowledge.PaginationInfo; +import software.amazon.smithy.model.shapes.OperationShape; +import software.amazon.smithy.model.shapes.ServiceShape; +import software.amazon.smithy.model.shapes.Shape; +import software.amazon.smithy.model.traits.PaginatedTrait; + +/** + * Flattens pagination info from service shapes into operation-level pagination traits. + */ +final class FlattenPaginationInfo { + + private final ServiceShape service; + + FlattenPaginationInfo(ServiceShape service) { + this.service = service; + } + + public Model transform(ModelTransformer transformer, Model model) { + Optional serviceLevelPagination = service.getTrait(PaginatedTrait.class); + if (!serviceLevelPagination.isPresent()) { + return model; + } + PaginatedIndex paginatedIndex = PaginatedIndex.of(model); + + // Merge service-level information into each operation's pagination trait. + Set updatedShapes = new HashSet<>(); + for (OperationShape operationShape : model.getOperationShapesWithTrait(PaginatedTrait.class)) { + PaginationInfo paginationInfo = paginatedIndex.getPaginationInfo(service, operationShape).get(); + OperationShape updatedShape = operationShape.toBuilder() + .addTrait(paginationInfo.getPaginatedTrait()) + .build(); + updatedShapes.add(updatedShape); + } + + // Remove the paginated trait from the service as it's info has been flattened into the operations + updatedShapes.add(service.toBuilder().removeTrait(PaginatedTrait.ID).build()); + + return transformer.replaceShapes(model, updatedShapes); + } +} diff --git a/smithy-model/src/main/java/software/amazon/smithy/model/transform/ModelTransformer.java b/smithy-model/src/main/java/software/amazon/smithy/model/transform/ModelTransformer.java index 5d5ea4769c3..4c6160db4a6 100644 --- a/smithy-model/src/main/java/software/amazon/smithy/model/transform/ModelTransformer.java +++ b/smithy-model/src/main/java/software/amazon/smithy/model/transform/ModelTransformer.java @@ -677,4 +677,14 @@ public Model removeInvalidDefaults(Model model) { public Model deconflictErrorsWithSharedStatusCode(Model model, ServiceShape forService) { return new DeconflictErrorsWithSharedStatusCode(forService).transform(this, model); } + + /** + * Flattens all service-level pagination information into operation-level pagination traits. + * + * @param model Model to transform. + * @return Returns the transformed model. + */ + public Model flattenPaginationInfoIntoOperations(Model model, ServiceShape forService) { + return new FlattenPaginationInfo(forService).transform(this, model); + } } diff --git a/smithy-model/src/test/java/software/amazon/smithy/model/transform/FlattenPaginationInfoTest.java b/smithy-model/src/test/java/software/amazon/smithy/model/transform/FlattenPaginationInfoTest.java new file mode 100644 index 00000000000..3b35ebb5744 --- /dev/null +++ b/smithy-model/src/test/java/software/amazon/smithy/model/transform/FlattenPaginationInfoTest.java @@ -0,0 +1,42 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.model.transform; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; + +import org.junit.jupiter.api.Test; +import software.amazon.smithy.model.Model; +import software.amazon.smithy.model.shapes.ServiceShape; +import software.amazon.smithy.model.shapes.Shape; +import software.amazon.smithy.model.shapes.ShapeId; +import software.amazon.smithy.model.traits.PaginatedTrait; + +public class FlattenPaginationInfoTest { + private static final ShapeId serviceId = ShapeId.from("smithy.example#PaginatedService"); + private static final ShapeId operationId = ShapeId.from("smithy.example#PaginatedOperation"); + + @Test + void compareTransform() { + Model before = Model.assembler() + .addImport(FlattenPaginationInfoTest.class.getResource("flatten-pagination-before.smithy")) + .assemble() + .unwrap(); + ServiceShape service = before.expectShape(serviceId).asServiceShape().get(); + Model result = ModelTransformer.create().flattenPaginationInfoIntoOperations(before, service); + + Shape resultService = result.expectShape(serviceId); + assertFalse(resultService.hasTrait(PaginatedTrait.class)); + + + Shape resultOperation = result.expectShape(operationId); + PaginatedTrait resultTrait = resultOperation.expectTrait(PaginatedTrait.class); + assertEquals(resultTrait.getInputToken().get(), "nextToken"); + assertEquals(resultTrait.getOutputToken().get(), "nextToken"); + assertEquals(resultTrait.getPageSize().get(), "maxResults"); + assertEquals(resultTrait.getItems().get(), "foos"); + } +} diff --git a/smithy-model/src/test/resources/software/amazon/smithy/model/transform/flatten-pagination-before.smithy b/smithy-model/src/test/resources/software/amazon/smithy/model/transform/flatten-pagination-before.smithy new file mode 100644 index 00000000000..1534b0706a3 --- /dev/null +++ b/smithy-model/src/test/resources/software/amazon/smithy/model/transform/flatten-pagination-before.smithy @@ -0,0 +1,29 @@ +$version: "2.0" + +namespace smithy.example + +@paginated(inputToken: "nextToken", outputToken: "nextToken") +service PaginatedService { + operations: [ + PaginatedOperation + ] +} + +@paginated(pageSize: "maxResults", items: "foos") +operation PaginatedOperation { + input := { + maxResults: Integer + nextToken: String + } + output := { + nextToken: String + + @required + foos: StringList + } +} + +@private +list StringList { + member: String +}