Add customization to S3 ListParts Paginator

Adding new trait IsTruncatedPaginatorTrait that indicates the Paginator should use the value of the isTruncated field to indicate if the pagination has been exhausted.
smithy-lang · Jun 3, 2024 · 6816dd8 · 6816dd8
1 parent 9c1ae5a
commit 6816dd8
Show file tree

Hide file tree

Showing 3 changed files with 76 additions and 3 deletions.
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
@@ -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/...kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/PaginatorGenerator.kt b/...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<SyntheticOutputTrait>()?.originalId.let { shapeId ->
+                codegenContext.model.getShape(shapeId).orNull()?.hasTrait<IsTruncatedPaginatorTrait>() ?: 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/...lin/software/amazon/smithy/rust/codegen/client/smithy/traits/IsTruncatedPaginatorTrait.kt b/...lin/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")
+    }
+}