From 98f94961dd8d8461da2ee39897e03543ce717a0b Mon Sep 17 00:00:00 2001
From: David Yaffe <dayaffe@amazon.com>
Date: Thu, 30 Nov 2023 15:39:45 -0500
Subject: [PATCH] add IsTruncated pagination customization for S3 ListParts

---
 .../s3/TruncatablePaginationIntegration.kt    | 39 +++++++++++++++++++
 ...swift.codegen.integration.SwiftIntegration |  1 +
 2 files changed, 40 insertions(+)
 create mode 100644 codegen/smithy-aws-swift-codegen/src/main/kotlin/software/amazon/smithy/aws/swift/codegen/customization/s3/TruncatablePaginationIntegration.kt

diff --git a/codegen/smithy-aws-swift-codegen/src/main/kotlin/software/amazon/smithy/aws/swift/codegen/customization/s3/TruncatablePaginationIntegration.kt b/codegen/smithy-aws-swift-codegen/src/main/kotlin/software/amazon/smithy/aws/swift/codegen/customization/s3/TruncatablePaginationIntegration.kt
new file mode 100644
index 00000000000..d3562f906b7
--- /dev/null
+++ b/codegen/smithy-aws-swift-codegen/src/main/kotlin/software/amazon/smithy/aws/swift/codegen/customization/s3/TruncatablePaginationIntegration.kt
@@ -0,0 +1,39 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package software.amazon.smithy.aws.swift.codegen.customization.s3
+
+import software.amazon.smithy.model.Model
+import software.amazon.smithy.model.shapes.MemberShape
+import software.amazon.smithy.model.shapes.ServiceShape
+import software.amazon.smithy.model.transform.ModelTransformer
+import software.amazon.smithy.swift.codegen.SwiftSettings
+import software.amazon.smithy.swift.codegen.customtraits.PaginationTruncationMember
+import software.amazon.smithy.swift.codegen.integration.SwiftIntegration
+import software.amazon.smithy.swift.codegen.model.expectShape
+
+private val TRUNCATION_MEMBER_IDS = setOf(
+    "com.amazonaws.s3#ListPartsOutput\$IsTruncated",
+)
+
+/**
+ * Applies the [PaginationTruncationMember] annotation to a manually-curated list of operations and members to handle
+ * non-standard pagination termination conditions.
+ */
+class TruncatablePaginationIntegration : SwiftIntegration {
+    override fun enabledForService(model: Model, settings: SwiftSettings): Boolean =
+        model.expectShape<ServiceShape>(settings.service).isS3
+
+    override fun preprocessModel(model: Model, settings: SwiftSettings): Model = ModelTransformer
+        .create()
+        .mapShapes(model) { shape ->
+            when {
+                shape.id.toString() in TRUNCATION_MEMBER_IDS -> {
+                    check(shape is MemberShape) { "Cannot apply PaginationTruncationMember to non-member shape" }
+                    shape.toBuilder().addTrait(PaginationTruncationMember()).build()
+                }
+                else -> shape
+            }
+        }
+}
diff --git a/codegen/smithy-aws-swift-codegen/src/main/resources/META-INF/services/software.amazon.smithy.swift.codegen.integration.SwiftIntegration b/codegen/smithy-aws-swift-codegen/src/main/resources/META-INF/services/software.amazon.smithy.swift.codegen.integration.SwiftIntegration
index ea2f756d8d1..e6566d3ed91 100644
--- a/codegen/smithy-aws-swift-codegen/src/main/resources/META-INF/services/software.amazon.smithy.swift.codegen.integration.SwiftIntegration
+++ b/codegen/smithy-aws-swift-codegen/src/main/resources/META-INF/services/software.amazon.smithy.swift.codegen.integration.SwiftIntegration
@@ -1,6 +1,7 @@
 software.amazon.smithy.aws.swift.codegen.AddProtocols
 software.amazon.smithy.aws.swift.codegen.customization.s3.S3ErrorIntegration
 software.amazon.smithy.aws.swift.codegen.customization.s3.S3Expires
+software.amazon.smithy.aws.swift.codegen.customization.s3.TruncatablePaginationIntegration
 software.amazon.smithy.aws.swift.codegen.customization.route53.Route53TrimHostedZone
 software.amazon.smithy.aws.swift.codegen.customization.route53.Route53InvalidBatchErrorIntegration
 software.amazon.smithy.aws.swift.codegen.customization.apigateway.ApiGatewayAddAcceptHeader