diff --git a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/s3/S3Decorator.kt b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/s3/S3Decorator.kt index 6f3577fb982..4c90a28b8ca 100644 --- a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/s3/S3Decorator.kt +++ b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/s3/S3Decorator.kt @@ -28,6 +28,7 @@ import software.amazon.smithy.rust.codegen.client.smithy.generators.OperationGen import software.amazon.smithy.rust.codegen.client.smithy.generators.OperationSection import software.amazon.smithy.rust.codegen.client.smithy.protocols.ClientRestXmlFactory import software.amazon.smithy.rust.codegen.client.smithy.traits.IncompatibleWithStalledStreamProtectionTrait +import software.amazon.smithy.rust.codegen.client.smithy.traits.IsTruncatedPaginatorTrait import software.amazon.smithy.rust.codegen.core.rustlang.Writable import software.amazon.smithy.rust.codegen.core.rustlang.rustBlockTemplate import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate @@ -64,6 +65,7 @@ class S3Decorator : ClientCodegenDecorator { setOf( ShapeId.from("com.amazonaws.s3#CopyObject"), ) + private val operationsWithIsTruncatedPaginator = setOf(ShapeId.from("com.amazonaws.s3#ListPartsOutput")) override fun protocols( serviceId: ShapeId, @@ -89,6 +91,9 @@ class S3Decorator : ClientCodegenDecorator { }.letIf(operationsIncompatibleWithStalledStreamProtection.contains(shape.id)) { logger.info("Adding IncompatibleWithStalledStreamProtection trait to $it") (it as OperationShape).toBuilder().addTrait(IncompatibleWithStalledStreamProtectionTrait()).build() + }.letIf(isInIsTruncatedList(shape)) { + logger.info("Adding IsTruncatedPaginator trait to $it") + (it as StructureShape).toBuilder().addTrait(IsTruncatedPaginatorTrait()).build() } } // the model has the bucket in the path @@ -155,6 +160,7 @@ class S3Decorator : ClientCodegenDecorator { ) } } + else -> {} } } @@ -165,6 +171,10 @@ class S3Decorator : ClientCodegenDecorator { private fun isInInvalidXmlRootAllowList(shape: Shape): Boolean { return shape.isStructureShape && invalidXmlRootAllowList.contains(shape.id) } + + private fun isInIsTruncatedList(shape: Shape): Boolean { + return shape.isStructureShape && operationsWithIsTruncatedPaginator.contains(shape.id) + } } class FilterEndpointTests( diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/PaginatorGenerator.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/PaginatorGenerator.kt index 3d9700993f4..3c3921ecd75 100644 --- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/PaginatorGenerator.kt +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/PaginatorGenerator.kt @@ -11,6 +11,7 @@ import software.amazon.smithy.model.shapes.OperationShape import software.amazon.smithy.model.traits.IdempotencyTokenTrait import software.amazon.smithy.model.traits.PaginatedTrait import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext +import software.amazon.smithy.rust.codegen.client.smithy.traits.IsTruncatedPaginatorTrait import software.amazon.smithy.rust.codegen.core.rustlang.RustModule import software.amazon.smithy.rust.codegen.core.rustlang.RustType import software.amazon.smithy.rust.codegen.core.rustlang.Writable @@ -21,8 +22,10 @@ import software.amazon.smithy.rust.codegen.core.rustlang.writable import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType.Companion.preludeScope import software.amazon.smithy.rust.codegen.core.smithy.rustType +import software.amazon.smithy.rust.codegen.core.smithy.traits.SyntheticOutputTrait import software.amazon.smithy.rust.codegen.core.util.PANIC import software.amazon.smithy.rust.codegen.core.util.findMemberWithTrait +import software.amazon.smithy.rust.codegen.core.util.getTrait import software.amazon.smithy.rust.codegen.core.util.hasTrait import software.amazon.smithy.rust.codegen.core.util.inputShape import software.amazon.smithy.rust.codegen.core.util.orNull @@ -71,6 +74,13 @@ class PaginatorGenerator private constructor( private val outputType = symbolProvider.toSymbol(outputShape) private val errorType = symbolProvider.symbolForOperationError(operation) + private val isTruncatedPaginator = + codegenContext.model.getShape(outputShape.toShapeId()).orNull().let { shape -> + shape?.getTrait()?.originalId.let { shapeId -> + codegenContext.model.getShape(shapeId).orNull()?.hasTrait() ?: false + } + } + private fun paginatorType(): RuntimeType = RuntimeType.forInlineFun( paginatorName, @@ -89,7 +99,9 @@ class PaginatorGenerator private constructor( "Error" to errorType, "Builder" to symbolProvider.symbolForBuilder(operation.inputShape(model)), // SDK Types - "HttpResponse" to RuntimeType.smithyRuntimeApiClient(runtimeConfig).resolve("client::orchestrator::HttpResponse"), + "HttpResponse" to + RuntimeType.smithyRuntimeApiClient(runtimeConfig) + .resolve("client::orchestrator::HttpResponse"), "SdkError" to RuntimeType.sdkError(runtimeConfig), "pagination_stream" to RuntimeType.smithyAsync(runtimeConfig).resolve("future::pagination_stream"), // External Types @@ -161,7 +173,7 @@ class PaginatorGenerator private constructor( let done = match resp { #{Ok}(ref resp) => { let new_token = #{output_token}(resp); - let is_empty = new_token.map(|token| token.is_empty()).unwrap_or(true); + #{is_empty_setter:W} if !is_empty && new_token == input.$inputTokenMember.as_ref() && self.stop_on_duplicate_token { true } else { @@ -211,9 +223,37 @@ class PaginatorGenerator private constructor( "RuntimePlugins" to RuntimeType.runtimePlugins(runtimeConfig), ) }, + "is_empty_setter" to isEmptySetter(), ) } + /** Generate code to calculate the value of is_empty. For most paginators this + * is indicated by the next token being the empty string. But for paginators + * with the isTruncatedPaginator trait the next token is not necessarily empty. + * (ex: for s3 ListParts the final next token is "0" when pagination is complete, + * causing the paginator to go back to the first page and loop forever) + * In this case we use a false value of isTruncated as the only indicator that + * the pagination is exhausted. + * */ + private fun isEmptySetter() = + writable { + if (isTruncatedPaginator) { + rustTemplate( + """ + // Pagination is exhausted when `is_truncated` is false + let is_empty = !resp.is_truncated.unwrap_or(false); + """, + ) + } else { + rustTemplate( + """ + // Pagination is exhausted when the next token is an empty string + let is_empty = new_token.map(|token| token.is_empty()).unwrap_or(true); + """, + ) + } + } + /** Type of the inner item of the paginator */ private fun itemType(): String { val members = paginationInfo.itemsMemberPath @@ -280,7 +320,10 @@ class PaginatorGenerator private constructor( ), "item_type" to writable { - rustTemplate("#{Result}<${itemType()}, #{SdkError}<#{Error}, #{HttpResponse}>>", *codegenScope) + rustTemplate( + "#{Result}<${itemType()}, #{SdkError}<#{Error}, #{HttpResponse}>>", + *codegenScope, + ) }, *codegenScope, ) diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/traits/IsTruncatedPaginatorTrait.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/traits/IsTruncatedPaginatorTrait.kt new file mode 100644 index 00000000000..8487298e909 --- /dev/null +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/traits/IsTruncatedPaginatorTrait.kt @@ -0,0 +1,20 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.rust.codegen.client.smithy.traits + +import software.amazon.smithy.model.node.Node +import software.amazon.smithy.model.shapes.ShapeId +import software.amazon.smithy.model.traits.AnnotationTrait + +/** + * Indicates that an operation should use the IsTruncated field for detecting the end of pagination. + */ +class IsTruncatedPaginatorTrait : AnnotationTrait(ID, Node.objectNode()) { + companion object { + val ID: ShapeId = + ShapeId.from("software.amazon.smithy.rust.codegen.client.smithy.traits#isTruncatedPaginatorTrait") + } +}