feat: Paginator codegen using Kotlin Flow (#557)

.gitignore

@@ -33,6 +33,7 @@ hs_err_pid*
 # Gradle
 .gradle
 build/
+local.properties
 
 # Ignore Gradle GUI config
 gradle-app.setting

runtime/runtime-core/build.gradle.kts

@@ -23,7 +23,7 @@ kotlin {
                 // Attributes property bag is exposed as client options
                 api(project(":runtime:utils"))
                 // Coroutines' locking features are used in retry token bucket implementations
-                implementation("org.jetbrains.kotlinx:kotlinx-coroutines-core:$coroutinesVersion")
+                api("org.jetbrains.kotlinx:kotlinx-coroutines-core:$coroutinesVersion")
             }
         }
 

settings.gradle.kts

@@ -44,10 +44,11 @@ include(":runtime:protocol:http")
 include(":runtime:protocol:http-test")
 include(":runtime:protocol:http-client-engines:http-client-engine-ktor")
 
-include(":compile-tests")
-include(":benchmarks")
-include(":benchmarks:serde-benchmarks-codegen")
-include(":benchmarks:serde-benchmarks")
+include(":tests")
+include(":tests:benchmarks:serde-benchmarks-codegen")
+include(":tests:benchmarks:serde-benchmarks")
+include(":tests:compile")
+include(":tests:codegen:paginator-tests")
 
 include(":dokka-smithy")
 include(":ktlint-rules")

smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/core/ExternalTypes.kt

@@ -0,0 +1,19 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0.
+ */
+package software.amazon.smithy.kotlin.codegen.core
+
+import software.amazon.smithy.kotlin.codegen.model.toSymbol
+
+/**
+ * Commonly used external types.
+ */
+object ExternalTypes {
+    // https://github.com/Kotlin/kotlinx.coroutines
+    object KotlinxCoroutines {
+        val Flow = "kotlinx.coroutines.flow.Flow".toSymbol()
+        val FlowGenerator = "kotlinx.coroutines.flow.flow".toSymbol()
+        val FlowTransform = "kotlinx.coroutines.flow.transform".toSymbol()
+    }
+}

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

@@ -0,0 +1,237 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0.
+ */
+package software.amazon.smithy.kotlin.codegen.rendering
+
+import software.amazon.smithy.codegen.core.CodegenException
+import software.amazon.smithy.codegen.core.Symbol
+import software.amazon.smithy.codegen.core.SymbolReference
+import software.amazon.smithy.kotlin.codegen.KotlinSettings
+import software.amazon.smithy.kotlin.codegen.core.CodegenContext
+import software.amazon.smithy.kotlin.codegen.core.ExternalTypes
+import software.amazon.smithy.kotlin.codegen.core.KotlinDelegator
+import software.amazon.smithy.kotlin.codegen.core.KotlinWriter
+import software.amazon.smithy.kotlin.codegen.core.defaultName
+import software.amazon.smithy.kotlin.codegen.core.withBlock
+import software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
+import software.amazon.smithy.kotlin.codegen.model.SymbolProperty
+import software.amazon.smithy.kotlin.codegen.model.expectShape
+import software.amazon.smithy.kotlin.codegen.model.hasTrait
+import software.amazon.smithy.kotlin.codegen.utils.getOrNull
+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.CollectionShape
+import software.amazon.smithy.model.shapes.MapShape
+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
+
+/**
+ * Generate paginators for supporting operations.  See
+ * https://awslabs.github.io/smithy/1.0/spec/core/behavior-traits.html#paginated-trait for details.
+ */
+class PaginatorGenerator : KotlinIntegration {
+    override fun enabledForService(model: Model, settings: KotlinSettings): Boolean =
+        model.operationShapes.any { it.hasTrait<PaginatedTrait>() }
+
+    override fun writeAdditionalFiles(ctx: CodegenContext, delegator: KotlinDelegator) {
+        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) }
+
+            paginatedOperations.forEach { paginatedOperation ->
+                val paginationInfo = paginatedIndex.getPaginationInfo(service, paginatedOperation).getOrNull()
+                    ?: throw CodegenException("Unexpectedly unable to get PaginationInfo from $service $paginatedOperation")
+                val paginationItemInfo = getItemDescriptorOrNull(paginationInfo, ctx)
+
+                renderPaginatorForOperation(writer, ctx, service, paginatedOperation, paginationInfo, paginationItemInfo)
+            }
+        }
+    }
+
+    // Render paginator(s) for operation
+    private fun renderPaginatorForOperation(
+        writer: KotlinWriter,
+        ctx: CodegenContext,
+        service: ServiceShape,
+        paginatedOperation: OperationShape,
+        paginationInfo: PaginationInfo,
+        itemDesc: ItemDescriptor?
+    ) {
+        val serviceSymbol = ctx.symbolProvider.toSymbol(service)
+        val outputSymbol = ctx.symbolProvider.toSymbol(paginationInfo.output)
+        val inputSymbol = ctx.symbolProvider.toSymbol(paginationInfo.input)
+
+        renderResponsePaginator(
+            writer,
+            serviceSymbol,
+            paginatedOperation,
+            inputSymbol,
+            outputSymbol,
+            paginationInfo
+        )
+
+        // Optionally generate paginator when nested item is specified on the trait.
+        if (itemDesc != null) {
+            renderItemPaginator(
+                writer,
+                service,
+                paginatedOperation,
+                itemDesc,
+                outputSymbol
+            )
+        }
+    }
+
+    // Generate the paginator that iterates over responses
+    private fun renderResponsePaginator(
+        writer: KotlinWriter,
+        serviceSymbol: Symbol,
+        operationShape: OperationShape,
+        inputSymbol: Symbol,
+        outputSymbol: Symbol,
+        paginationInfo: PaginationInfo
+    ) {
+        val nextMarkerLiteral = paginationInfo.outputTokenMemberPath.joinToString(separator = "?.") {
+            it.defaultName()
+        }
+        val markerLiteral = paginationInfo.inputTokenMember.defaultName()
+
+        writer.write("")
+        writer.dokka(
+            """
+                  Paginate over [${outputSymbol.name}] 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 [${inputSymbol.name}] to start pagination
+                  @return A [kotlinx.coroutines.flow.Flow] that can collect [${outputSymbol.name}]
+            """.trimIndent()
+        )
+        writer
+            .addImport(ExternalTypes.KotlinxCoroutines.Flow)
+            .addImport(ExternalTypes.KotlinxCoroutines.FlowGenerator)
+            .addImport(serviceSymbol)
+            .addImport(inputSymbol)
+            .addImport(outputSymbol)
+            .withBlock(
+                "fun #T.#LPaginated(initialRequest: #T): Flow<#T> =",
+                "",
+                serviceSymbol,
+                operationShape.defaultName(),
+                inputSymbol,
+                outputSymbol
+            ) {
+                withBlock("flow {", "}") {
+                    write("var cursor: String? = null")
+                    write("var isFirstPage: Boolean = true")
+                    write("")
+                    withBlock("while (isFirstPage || (cursor?.isNotEmpty() == true)) {", "}") {
+                        withBlock("val req = initialRequest.copy {", "}") {
+                            write("this.$markerLiteral = cursor")
+                        }
+                        write(
+                            "val result = this@#1LPaginated.#1L(req)",
+                            operationShape.defaultName()
+                        )
+                        write("isFirstPage = false")
+                        write("cursor = result.$nextMarkerLiteral")
+                        write("emit(result)")
+                    }
+                }
+            }
+    }
+
+    // Generate a paginator that iterates over the model-specified item
+    private fun renderItemPaginator(
+        writer: KotlinWriter,
+        serviceShape: ServiceShape,
+        operationShape: OperationShape,
+        itemDesc: ItemDescriptor,
+        outputSymbol: Symbol,
+    ) {
+        writer.write("")
+        writer.dokka(
+            """
+                This paginator transforms the flow returned by [${operationShape.defaultName()}Paginated] 
+                to access the nested member [${itemDesc.targetMember.defaultName(serviceShape)}]
+                @return A [kotlinx.coroutines.flow.Flow] that can collect [${itemDesc.targetMember.defaultName(serviceShape)}]
+            """.trimIndent()
+        )
+        writer
+            .addImport(ExternalTypes.KotlinxCoroutines.FlowTransform)
+            .addImport(itemDesc.itemSymbol)
+            .addImportReferences(itemDesc.itemSymbol, SymbolReference.ContextOption.USE)
+            // @JvmName is required due to Java interop compatibility in the compiler.
+            // Multiple functions may have the same name and the generic does not disambiguate the type in Java.
+            // NOTE: This does not mean these functions are callable from Java.
+            .write(
+                """@JvmName("#L#L")""",
+                outputSymbol.name.replaceFirstChar(Char::lowercaseChar),
+                itemDesc.targetMember.defaultName(serviceShape)
+            )
+            .withBlock(
+                "fun #T<#T>.#L(): #T<#L> =", "",
+                ExternalTypes.KotlinxCoroutines.Flow,
+                outputSymbol,
+                itemDesc.itemLiteral,
+                ExternalTypes.KotlinxCoroutines.Flow,
+                itemDesc.collectionLiteral
+            ) {
+                withBlock("transform() { response -> ", "}") {
+                    withBlock("response.#L?.forEach {", "}", itemDesc.itemPathLiteral) {
+                        write("emit(it)")
+                    }
+                }
+            }
+    }
+}
+
+/**
+ * Model info necessary to codegen paginator item
+ */
+private data class ItemDescriptor(
+    val collectionLiteral: String,
+    val targetMember: Shape,
+    val itemLiteral: String,
+    val itemPathLiteral: String,
+    val itemSymbol: Symbol
+)
+
+/**
+ * Return an [ItemDescriptor] if model supplies, otherwise null
+ */
+private fun getItemDescriptorOrNull(paginationInfo: PaginationInfo, ctx: CodegenContext): ItemDescriptor? {
+    val itemMemberId = paginationInfo.itemsMemberPath?.lastOrNull()?.target ?: return null
+
+    val itemLiteral = paginationInfo.itemsMemberPath!!.last()!!.defaultName()
+    val itemPathLiteral = paginationInfo.itemsMemberPath.joinToString(separator = "?.") { it.defaultName() }
+    val itemMember = ctx.model.expectShape(itemMemberId)
+    val (collectionLiteral, targetMember) = when (itemMember) {
+        is MapShape ->
+            ctx.symbolProvider.toSymbol(itemMember)
+                .expectProperty(SymbolProperty.ENTRY_EXPRESSION) as String to itemMember
+        is CollectionShape ->
+            ctx.symbolProvider.toSymbol(ctx.model.expectShape(itemMember.member.target)).name to ctx.model.expectShape(
+                itemMember.member.target
+            )
+        else -> error("Unexpected shape type ${itemMember.type}")
+    }
+
+    return ItemDescriptor(
+        collectionLiteral,
+        targetMember,
+        itemLiteral,
+        itemPathLiteral,
+        ctx.symbolProvider.toSymbol(itemMember)
+    )
+}

smithy-kotlin-codegen/src/main/resources/META-INF/services/software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration

@@ -1,2 +1,3 @@
 software.amazon.smithy.kotlin.codegen.lang.BuiltinPreprocessor
 software.amazon.smithy.kotlin.codegen.lang.DocumentationPreprocessor
+software.amazon.smithy.kotlin.codegen.rendering.PaginatorGenerator