fix: correctly codegen waiters and paginators for resource operations (…

…#1064)

.changes/e26e1980-ee66-4d9d-99b5-6ee0f1f3901d.json

@@ -0,0 +1,8 @@
+{
+    "id": "e26e1980-ee66-4d9d-99b5-6ee0f1f3901d",
+    "type": "bugfix",
+    "description": "Correctly generate waiters and paginators for resource operations",
+    "issues": [
+        "awslabs/aws-sdk-kotlin#900"
+    ]
+}

codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/PaginatorGenerator.kt

@@ -47,11 +47,12 @@ class PaginatorGenerator : KotlinIntegration {
         val service = ctx.model.expectShape<ServiceShape>(ctx.settings.service)
         val paginatedIndex = PaginatedIndex.of(ctx.model)
 
-        delegator.useFileWriter("Paginators.kt", "${ctx.settings.pkg.name}.paginators") { writer ->
-            val paginatedOperations = service.allOperations
-                .map { ctx.model.expectShape<OperationShape>(it) }
-                .filter { operationShape -> operationShape.hasTrait(PaginatedTrait.ID) }
+        val paginatedOperations = TopDownIndex
+            .of(ctx.model)
+            .getContainedOperations(ctx.settings.service)
+            .filter { it.hasTrait<PaginatedTrait>() }
 
+        delegator.useFileWriter("Paginators.kt", "${ctx.settings.pkg.name}.paginators") { writer ->
             paginatedOperations.forEach { paginatedOperation ->
                 val paginationInfo = paginatedIndex.getPaginationInfo(service, paginatedOperation).getOrNull()
                     ?: throw CodegenException("Unexpectedly unable to get PaginationInfo from $service $paginatedOperation")

codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/waiters/ServiceWaitersGenerator.kt

@@ -44,9 +44,9 @@ internal fun CodegenContext.allWaiters(): List<WaiterInfo> {
             WaiterInfo(this, service, op, name, waiter)
         } ?: listOf()
 
-    return service
-        .allOperations
-        .map { model.expectShape<OperationShape>(it) }
+    return TopDownIndex
+        .of(model)
+        .getContainedOperations(service)
         .flatMap(::operationWaiters)
 }
 

codegen/smithy-kotlin-codegen/src/test/kotlin/software/amazon/smithy/kotlin/codegen/rendering/PaginatorGeneratorTest.kt

@@ -424,4 +424,116 @@ class PaginatorGeneratorTest {
 
         actual.shouldContainOnlyOnceWithDiff(expected)
     }
+
+    @Test
+    fun testRenderPaginatorFromResourceOperation() {
+        val testModel = """
+            namespace com.test
+            
+            use aws.protocols#restJson1
+            
+            service Lambda {
+                resources: [Function],
+            }
+            
+            resource Function {
+                identifiers: { id: String },
+                list: ListFunctions,
+            }
+            
+            @paginated(
+                inputToken: "Marker",
+                outputToken: "NextMarker",
+                pageSize: "MaxItems",
+            )
+            @readonly
+            @http(method: "GET", uri: "/functions", code: 200)
+            operation ListFunctions {
+                input: ListFunctionsRequest,
+                output: ListFunctionsResponse,
+            }
+            
+            structure ListFunctionsRequest {
+                @httpQuery("Marker")
+                Marker: String,
+                @httpQuery("MaxItems")
+                MaxItems: Integer,
+            }
+            
+            structure ListFunctionsResponse {
+                Functions: FunctionsList,
+                NextMarker: String,
+            }
+            
+            list FunctionsList {
+                member: FunctionSummary,
+            }
+            
+            structure FunctionSummary {
+                id: String,
+                name: String
+            }
+        """.toSmithyModel()
+        val testContext = testModel.newTestContext("Lambda", "com.test")
+
+        val codegenContext = object : CodegenContext {
+            override val model: Model = testContext.generationCtx.model
+            override val symbolProvider: SymbolProvider = testContext.generationCtx.symbolProvider
+            override val settings: KotlinSettings = testContext.generationCtx.settings
+            override val protocolGenerator: ProtocolGenerator = testContext.generator
+            override val integrations: List<KotlinIntegration> = testContext.generationCtx.integrations
+        }
+
+        val unit = PaginatorGenerator()
+        unit.writeAdditionalFiles(codegenContext, testContext.generationCtx.delegator)
+
+        testContext.generationCtx.delegator.flushWriters()
+        val testManifest = testContext.generationCtx.delegator.fileManifest as MockManifest
+        val actual = testManifest.expectFileString("src/main/kotlin/com/test/paginators/Paginators.kt")
+
+        val expected = """
+            /**
+             * Paginate over [ListFunctionsResponse] results.
+             *
+             * When this operation is called, a [kotlinx.coroutines.Flow] is created. Flows are lazy (cold) so no service
+             * calls are made until the flow is collected. This also means there is no guarantee that the request is valid
+             * until then. Once you start collecting the flow, the SDK will lazily load response pages by making service
+             * calls until there are no pages left or the flow is cancelled. If there are errors in your request, you will
+             * see the failures only after you start collection.
+             * @param initialRequest A [ListFunctionsRequest] to start pagination
+             * @return A [kotlinx.coroutines.flow.Flow] that can collect [ListFunctionsResponse]
+             */
+            public fun TestClient.listFunctionsPaginated(initialRequest: ListFunctionsRequest = ListFunctionsRequest { }): Flow<ListFunctionsResponse> =
+                flow {
+                    var cursor: kotlin.String? = null
+                    var hasNextPage: Boolean = true
+            
+                    while (hasNextPage) {
+                        val req = initialRequest.copy {
+                            this.marker = cursor
+                        }
+                        val result = this@listFunctionsPaginated.listFunctions(req)
+                        cursor = result.nextMarker
+                        hasNextPage = cursor?.isNotEmpty() == true
+                        emit(result)
+                    }
+                }
+            
+            /**
+             * Paginate over [ListFunctionsResponse] results.
+             *
+             * When this operation is called, a [kotlinx.coroutines.Flow] is created. Flows are lazy (cold) so no service
+             * calls are made until the flow is collected. This also means there is no guarantee that the request is valid
+             * until then. Once you start collecting the flow, the SDK will lazily load response pages by making service
+             * calls until there are no pages left or the flow is cancelled. If there are errors in your request, you will
+             * see the failures only after you start collection.
+             * @param block A builder block used for DSL-style invocation of the operation
+             * @return A [kotlinx.coroutines.flow.Flow] that can collect [ListFunctionsResponse]
+             */
+            public fun TestClient.listFunctionsPaginated(block: ListFunctionsRequest.Builder.() -> Unit): Flow<ListFunctionsResponse> =
+                listFunctionsPaginated(ListFunctionsRequest.Builder().apply(block).build())
+        """.trimIndent()
+
+        actual.shouldContainOnlyOnceWithDiff(expected)
+    }
 }

codegen/smithy-kotlin-codegen/src/test/kotlin/software/amazon/smithy/kotlin/codegen/rendering/waiters/ServiceWaitersGeneratorTest.kt

@@ -20,37 +20,37 @@ import software.amazon.smithy.kotlin.codegen.test.*
 import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.shapes.ShapeId
 import kotlin.test.Test
-import kotlin.test.assertFalse
-import kotlin.test.assertTrue
+import kotlin.test.assertEquals
 
 class ServiceWaitersGeneratorTest {
-    private val generated = generateService("simple-service-with-waiter.smithy")
-
     @Test
     fun testServiceGate() {
-        val enabledServiceModel = loadModelFromResource("simple-service-with-waiter.smithy")
-        val enabledServiceSettings = enabledServiceModel.newTestContext().generationCtx.settings
-        assertTrue(ServiceWaitersGenerator().enabledForService(enabledServiceModel, enabledServiceSettings))
-
-        val disabledServiceModel = loadModelFromResource("simple-service.smithy")
-        val disabledServiceSettings = disabledServiceModel.newTestContext().generationCtx.settings
-        assertFalse(ServiceWaitersGenerator().enabledForService(disabledServiceModel, disabledServiceSettings))
+        mapOf(
+            "simple-service-with-operation-waiter.smithy" to true,
+            "simple-service-with-resource-waiter.smithy" to true,
+            "simple-service.smithy" to false,
+        ).forEach { (modelName, expectEnabled) ->
+            val model = loadModelFromResource(modelName)
+            val settings = model.newTestContext().generationCtx.settings
+            val actualEnabled = ServiceWaitersGenerator().enabledForService(model, settings)
+            assertEquals(expectEnabled, actualEnabled)
+        }
     }
 
     @Test
-    fun testMainWaiterMethod() {
+    fun testWaiterSignatureWithOptionalInput() {
         val methodHeader = """
             /**
-             * Wait until a foo exists
+             * Wait until a foo exists with optional input
              */
-            public suspend fun TestClient.waitUntilFooExists(request: DescribeFooRequest = DescribeFooRequest { }): Outcome<DescribeFooResponse> {
+            public suspend fun TestClient.waitUntilFooOptionalExists(request: DescribeFooOptionalRequest = DescribeFooOptionalRequest { }): Outcome<DescribeFooOptionalResponse> {
         """.trimIndent()
         val methodFooter = """
                 val policy = AcceptorRetryPolicy(request, acceptors)
-                return strategy.retry(policy) { describeFoo(request) }
+                return strategy.retry(policy) { describeFooOptional(request) }
             }
         """.trimIndent()
-        generated.shouldContain(methodHeader, methodFooter)
+        generateService("simple-service-with-operation-waiter.smithy").shouldContain(methodHeader, methodFooter)
     }
 
     @Test
@@ -61,30 +61,45 @@ class ServiceWaitersGeneratorTest {
              */
             public suspend fun TestClient.waitUntilFooRequiredExists(request: DescribeFooRequiredRequest): Outcome<DescribeFooRequiredResponse> {
         """.trimIndent()
-        generated.shouldContainOnlyOnceWithDiff(methodHeader)
+        listOf(
+            generateService("simple-service-with-operation-waiter.smithy"),
+            generateService("simple-service-with-resource-waiter.smithy"),
+        ).forEach { generated ->
+            generated.shouldContain(methodHeader)
+        }
     }
 
     @Test
     fun testConvenienceWaiterMethod() {
         val expected = """
             /**
-             * Wait until a foo exists
+             * Wait until a foo exists with required input
              */
-            public suspend fun TestClient.waitUntilFooExists(block: DescribeFooRequest.Builder.() -> Unit): Outcome<DescribeFooResponse> =
-                waitUntilFooExists(DescribeFooRequest.Builder().apply(block).build())
+            public suspend fun TestClient.waitUntilFooRequiredExists(block: DescribeFooRequiredRequest.Builder.() -> Unit): Outcome<DescribeFooRequiredResponse> =
+                waitUntilFooRequiredExists(DescribeFooRequiredRequest.Builder().apply(block).build())
         """.trimIndent()
-        generated.shouldContainOnlyOnce(expected)
+        listOf(
+            generateService("simple-service-with-operation-waiter.smithy"),
+            generateService("simple-service-with-resource-waiter.smithy"),
+        ).forEach { generated ->
+            generated.shouldContain(expected)
+        }
     }
 
     @Test
     fun testAcceptorList() {
         val expected = """
-            val acceptors = listOf<Acceptor<DescribeFooRequest, DescribeFooResponse>>(
+            val acceptors = listOf<Acceptor<DescribeFooRequiredRequest, DescribeFooRequiredResponse>>(
                 SuccessAcceptor(RetryDirective.TerminateAndSucceed, true),
                 ErrorTypeAcceptor(RetryDirective.RetryError(RetryErrorType.ServerSide), "NotFound"),
             )
         """.formatForTest()
-        generated.shouldContainOnlyOnce(expected)
+        listOf(
+            generateService("simple-service-with-operation-waiter.smithy"),
+            generateService("simple-service-with-resource-waiter.smithy"),
+        ).forEach { generated ->
+            generated.shouldContainOnlyOnce(expected)
+        }
     }
 
     @Test
@@ -101,7 +116,12 @@ class ServiceWaitersGeneratorTest {
                 }
             }
         """.formatForTest()
-        generated.shouldContain(expected)
+        listOf(
+            generateService("simple-service-with-operation-waiter.smithy"),
+            generateService("simple-service-with-resource-waiter.smithy"),
+        ).forEach { generated ->
+            generated.shouldContain(expected)
+        }
     }
 
     private fun generateService(modelResourceName: String): String {

codegen/smithy-kotlin-codegen/src/test/resources/software/amazon/smithy/kotlin/codegen/simple-service-with-waiter.smithy → codegen/smithy-kotlin-codegen/src/test/resources/software/amazon/smithy/kotlin/codegen/simple-service-with-operation-waiter.smithy

@@ -6,14 +6,14 @@ use smithy.waiters#waitable
 service Test {
     version: "1.0.0",
     operations: [
-        DescribeFoo,
+        DescribeFooOptional,
         DescribeFooRequired,
     ]
 }
 
 @waitable(
-    FooExists: {
-        documentation: "Wait until a foo exists",
+    FooOptionalExists: {
+        documentation: "Wait until a foo exists with optional input",
         acceptors: [
             {
                 state: "success",
@@ -30,8 +30,8 @@ service Test {
         ]
     }
 )
-operation DescribeFoo {
-    input: DescribeFooInput,
+operation DescribeFooOptional {
+    input: DescribeFooOptionalInput,
     output: DescribeFooOutput,
     errors: [NotFound, UnknownError]
 }
@@ -61,7 +61,7 @@ operation DescribeFooRequired {
     errors: [NotFound, UnknownError]
 }
 
-structure DescribeFooInput {
+structure DescribeFooOptionalInput {
     id: String
 }
 

codegen/smithy-kotlin-codegen/src/test/resources/software/amazon/smithy/kotlin/codegen/simple-service-with-resource-waiter.smithy

@@ -0,0 +1,57 @@
+$version: "1.0"
+namespace com.test
+
+use smithy.waiters#waitable
+
+service Test {
+    version: "1.0.0"
+    resources: [
+        Foo
+    ]
+}
+
+resource Foo {
+    identifiers: { id: String }
+    read: DescribeFooRequired
+}
+
+@readonly
+@waitable(
+    FooRequiredExists: {
+        documentation: "Wait until a foo exists with required input",
+        acceptors: [
+            {
+                state: "success",
+                matcher: {
+                    success: true
+                }
+            },
+            {
+                state: "retry",
+                matcher: {
+                    errorType: "NotFound"
+                }
+            }
+        ]
+    }
+)
+operation DescribeFooRequired {
+    input: DescribeFooRequiredInput,
+    output: DescribeFooOutput,
+    errors: [NotFound, UnknownError]
+}
+
+structure DescribeFooRequiredInput {
+    @required
+    id: String
+}
+
+structure DescribeFooOutput {
+    name: String
+}
+
+@error("client")
+structure NotFound {}
+
+@error("server")
+structure UnknownError {}