Skip to content

Commit cc8f42c

Browse files
ianbotsfianbotsf
authored
fix: add missing support for pagination which terminates by repeating the token (#1140)
1 parent 26e7536 commit cc8f42c

File tree

9 files changed

+232
-45
lines changed

9 files changed

+232
-45
lines changed

‎.changes/5eea083c-b810-4c6d-b636-75baa07b0c10.json‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
{
2+
"id": "5eea083c-b810-4c6d-b636-75baa07b0c10",
3+
"type": "bugfix",
4+
"description": "Add missing support for pagination which terminates by repeating the token (instead of returning no token)",
5+
"issues": [
6+
"awslabs/aws-sdk-kotlin#1326"
7+
]
8+
}

‎codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/model/traits/PaginationEndBehaviorTrait.kt‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
package software.amazon.smithy.kotlin.codegen.model.traits
6+
7+
import software.amazon.smithy.model.FromSourceLocation
8+
import software.amazon.smithy.model.SourceLocation
9+
import software.amazon.smithy.model.shapes.ShapeId
10+
import software.amazon.smithy.model.traits.StringTrait
11+
12+
class PaginationEndBehaviorTrait(
13+
val value: PaginationEndBehavior,
14+
sourceLocation: FromSourceLocation,
15+
) : StringTrait(ID, value.toString(), sourceLocation) {
16+
companion object {
17+
val ID = ShapeId.from("smithy.kotlin.traits#paginationEndBehavior")
18+
}
19+
20+
constructor(value: PaginationEndBehavior = PaginationEndBehavior.Default) : this(value, SourceLocation.NONE)
21+
22+
constructor(stringValue: String, sourceLocation: FromSourceLocation) :
23+
this(PaginationEndBehavior.fromString(stringValue), sourceLocation)
24+
25+
class Provider : StringTrait.Provider<PaginationEndBehaviorTrait>(ID, ::PaginationEndBehaviorTrait)
26+
}
27+
28+
sealed interface PaginationEndBehavior {
29+
data object OutputTokenEmpty : PaginationEndBehavior
30+
data object IdenticalToken : PaginationEndBehavior
31+
data class TruncationMember(val memberName: String) : PaginationEndBehavior
32+
33+
companion object {
34+
val Default = OutputTokenEmpty
35+
36+
fun fromString(stringValue: String): PaginationEndBehavior {
37+
val tokens = stringValue.split(":")
38+
return when (tokens[0]) {
39+
"OutputTokenEmpty" -> OutputTokenEmpty
40+
"IdenticalToken" -> IdenticalToken
41+
"TruncationMember" -> TruncationMember(tokens[1])
42+
else -> error("""Unknown PaginationEndBehavior type "${tokens[0]}"""")
43+
}
44+
}
45+
}
46+
}

‎codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/model/traits/PaginationTruncationMember.kt‎

Lines changed: 0 additions & 22 deletions
This file was deleted.

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

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,8 @@ import software.amazon.smithy.kotlin.codegen.core.withBlock
1717
import software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
1818
import software.amazon.smithy.kotlin.codegen.lang.KotlinTypes
1919
import software.amazon.smithy.kotlin.codegen.model.*
20-
import software.amazon.smithy.kotlin.codegen.model.traits.PaginationTruncationMember
20+
import software.amazon.smithy.kotlin.codegen.model.traits.PaginationEndBehavior
21+
import software.amazon.smithy.kotlin.codegen.model.traits.PaginationEndBehaviorTrait
2122
import software.amazon.smithy.kotlin.codegen.utils.getOrNull
2223
import software.amazon.smithy.model.Model
2324
import software.amazon.smithy.model.knowledge.PaginatedIndex
@@ -156,14 +157,19 @@ class PaginatorGenerator : KotlinIntegration {
156157
)
157158
write("cursor = result.$nextMarkerLiteral")
158159

159-
val hasNextPageFlag = outputShape
160-
.members()
161-
.singleOrNull { it.hasTrait(PaginationTruncationMember.ID) }
162-
?.defaultName()
163-
?.let { "result.$it" }
164-
?: "cursor?.isNotEmpty()"
160+
val endBehavior = operationShape.getTrait<PaginationEndBehaviorTrait>()?.value
161+
?: PaginationEndBehavior.Default
165162

166-
write("hasNextPage = #L == true", hasNextPageFlag)
163+
val hasNextPageFlag = when (endBehavior) {
164+
PaginationEndBehavior.OutputTokenEmpty -> "cursor?.isNotEmpty() == true"
165+
PaginationEndBehavior.IdenticalToken -> "cursor != null && cursor != req.$markerLiteral"
166+
is PaginationEndBehavior.TruncationMember -> {
167+
val member = outputShape.allMembers.getValue(endBehavior.memberName).defaultName()
168+
"result.$member == true" // $member will be a boolean flag
169+
}
170+
}
171+
172+
write("hasNextPage = #L", hasNextPageFlag)
167173
write("emit(result)")
168174
}
169175
}

‎codegen/smithy-kotlin-codegen/src/main/resources/META-INF/services/software.amazon.smithy.model.traits.TraitService‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
software.amazon.smithy.kotlin.codegen.model.traits.PaginationEndBehaviorTrait$Provider

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

Lines changed: 91 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -408,7 +408,7 @@ class PaginatorGeneratorTest {
408408
use aws.protocols#restJson1
409409
410410
@trait(selector: "*")
411-
structure paginationTruncationMember { }
411+
string paginationEndBehavior
412412
413413
service Lambda {
414414
operations: [ListFunctions]
@@ -419,6 +419,7 @@ class PaginatorGeneratorTest {
419419
outputToken: "NextMarker",
420420
pageSize: "MaxItems"
421421
)
422+
@paginationEndBehavior("TruncationMember:IsTruncated")
422423
@readonly
423424
@http(method: "GET", uri: "/functions", code: 200)
424425
operation ListFunctions {
@@ -439,7 +440,6 @@ class PaginatorGeneratorTest {
439440
440441
structure ListFunctionsResponse {
441442
Functions: FunctionConfigurationList,
442-
@paginationTruncationMember
443443
IsTruncated: Boolean,
444444
NextMarker: String
445445
}
@@ -490,6 +490,95 @@ class PaginatorGeneratorTest {
490490
actual.shouldContainOnlyOnceWithDiff(expectedCode)
491491
}
492492

493+
@Test
494+
fun testRenderPaginatorWithIdenticalTokenTerminator() {
495+
val testModel = """
496+
namespace smithy.kotlin.traits
497+
498+
use aws.protocols#restJson1
499+
500+
@trait(selector: "*")
501+
string paginationEndBehavior
502+
503+
service Lambda {
504+
operations: [ListFunctions]
505+
}
506+
507+
@paginated(
508+
inputToken: "Marker",
509+
outputToken: "NextMarker",
510+
pageSize: "MaxItems"
511+
)
512+
@paginationEndBehavior("IdenticalToken")
513+
@readonly
514+
@http(method: "GET", uri: "/functions", code: 200)
515+
operation ListFunctions {
516+
input: ListFunctionsRequest,
517+
output: ListFunctionsResponse
518+
}
519+
520+
structure ListFunctionsRequest {
521+
@httpQuery("FunctionVersion")
522+
FunctionVersion: String,
523+
@httpQuery("Marker")
524+
Marker: String,
525+
@httpQuery("MasterRegion")
526+
MasterRegion: String,
527+
@httpQuery("MaxItems")
528+
MaxItems: Integer,
529+
}
530+
531+
structure ListFunctionsResponse {
532+
Functions: FunctionConfigurationList,
533+
NextMarker: String
534+
}
535+
536+
list FunctionConfigurationList {
537+
member: FunctionConfiguration
538+
}
539+
540+
structure FunctionConfiguration {
541+
FunctionName: String
542+
}
543+
""".toSmithyModel()
544+
val testContext = testModel.newTestContext("Lambda", "smithy.kotlin.traits")
545+
546+
val codegenContext = object : CodegenContext {
547+
override val model = testContext.generationCtx.model
548+
override val symbolProvider = testContext.generationCtx.symbolProvider
549+
override val settings = testContext.generationCtx.settings
550+
override val protocolGenerator = testContext.generator
551+
override val integrations = testContext.generationCtx.integrations
552+
}
553+
554+
val unit = PaginatorGenerator()
555+
unit.writeAdditionalFiles(codegenContext, testContext.generationCtx.delegator)
556+
557+
testContext.generationCtx.delegator.flushWriters()
558+
val testManifest = testContext.generationCtx.delegator.fileManifest as MockManifest
559+
val actual = testManifest.expectFileString("src/main/kotlin/smithy/kotlin/traits/paginators/Paginators.kt")
560+
561+
val expectedCode = """
562+
public fun TestClient.listFunctionsPaginated(initialRequest: ListFunctionsRequest = ListFunctionsRequest { }): Flow<ListFunctionsResponse> =
563+
flow {
564+
var cursor: kotlin.String? = initialRequest.marker
565+
var hasNextPage: Boolean = true
566+
567+
while (hasNextPage) {
568+
val req = initialRequest.copy {
569+
this.marker = cursor
570+
}
571+
val result = [email protected](req)
572+
cursor = result.nextMarker
573+
hasNextPage = cursor != null && cursor != req.marker
574+
emit(result)
575+
}
576+
}
577+
""".trimIndent()
578+
579+
actual.shouldContainOnlyOnceWithDiff(expectedCode)
580+
}
581+
493582
@Test
494583
fun testRenderPaginatorWithRequiredInputMembers() {
495584
val testModelNoItem = """

‎tests/codegen/paginator-tests/build.gradle.kts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,9 @@ tasks.test {
4545
testLogging {
4646
events("passed", "skipped", "failed")
4747
showStandardStreams = true
48+
showStackTraces = true
49+
showExceptions = true
50+
exceptionFormat = org.gradle.api.tasks.testing.logging.TestExceptionFormat.FULL
4851
}
4952
}
5053

‎tests/codegen/paginator-tests/model/paginated-operations.smithy‎

Lines changed: 41 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,14 @@ namespace smithy.kotlin.traits
33
use aws.protocols#restJson1
44

55
@trait(selector: "*")
6-
structure paginationTruncationMember { }
6+
string paginationEndBehavior
77

88
service Lambda {
9-
operations: [ListFunctions, TruncatedListFunctions]
9+
operations: [ListFunctions, TruncatedListFunctions, IdenticalTokenListFunctions]
1010
}
1111

12+
// ListFunctions shapes
13+
1214
@paginated(
1315
inputToken: "Marker",
1416
outputToken: "NextMarker",
@@ -38,12 +40,15 @@ structure ListFunctionsResponse {
3840
NextMarker: String
3941
}
4042

43+
// TruncatedListFunctions shapes
44+
4145
@paginated(
4246
inputToken: "Marker",
4347
outputToken: "NextMarker",
4448
pageSize: "MaxItems",
4549
items: "Functions"
4650
)
51+
@paginationEndBehavior("TruncationMember:IsTruncated")
4752
@readonly
4853
@http(method: "GET", uri: "/truncatedFunctions", code: 200)
4954
operation TruncatedListFunctions {
@@ -64,11 +69,44 @@ structure TruncatedListFunctionsRequest {
6469

6570
structure TruncatedListFunctionsResponse {
6671
Functions: FunctionConfigurationList,
67-
@paginationTruncationMember
6872
IsTruncated: Boolean,
6973
NextMarker: String
7074
}
7175

76+
// IdenticalTokenListFunctions shapes
77+
78+
@paginated(
79+
inputToken: "Marker",
80+
outputToken: "NextMarker",
81+
pageSize: "MaxItems",
82+
items: "Functions"
83+
)
84+
@paginationEndBehavior("IdenticalToken")
85+
@readonly
86+
@http(method: "GET", uri: "/identicalTokenFunctions", code: 200)
87+
operation IdenticalTokenListFunctions {
88+
input: IdenticalTokenListFunctionsRequest,
89+
output: IdenticalTokenListFunctionsResponse
90+
}
91+
92+
structure IdenticalTokenListFunctionsRequest {
93+
@httpQuery("FunctionVersion")
94+
FunctionVersion: String,
95+
@httpQuery("Marker")
96+
Marker: String,
97+
@httpQuery("MasterRegion")
98+
MasterRegion: String,
99+
@httpQuery("MaxItems")
100+
MaxItems: Integer
101+
}
102+
103+
structure IdenticalTokenListFunctionsResponse {
104+
Functions: FunctionConfigurationList,
105+
NextMarker: String
106+
}
107+
108+
// Common shapes
109+
72110
list FunctionConfigurationList {
73111
member: FunctionConfiguration
74112
}

‎tests/codegen/paginator-tests/src/main/kotlin/smithy/kotlin/traits/DefaultLambdaClient.kt‎

Lines changed: 28 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -33,23 +33,41 @@ class TestLambdaClient : LambdaClient {
3333
var exhaustedVal: String? = null
3434

3535
override suspend fun listFunctions(input: ListFunctionsRequest) = ListFunctionsResponse {
36-
nextMarker = when {
37-
input.marker.toIntOrZero() == pageCount - 1 -> exhaustedVal // Exhausted pages
38-
input.marker == null -> "1" // First page
39-
else -> (input.marker.toInt() + 1).toString() // Next page
36+
val inputMarker = input.marker.toIntOrZero()
37+
38+
nextMarker = when (inputMarker) {
39+
pageCount - 1 -> exhaustedVal // Exhausted pages
40+
else -> (inputMarker + 1).toString() // Next page
4041
}
4142

42-
functions = generateFunctions(input.marker.toIntOrZero())
43+
functions = generateFunctions(inputMarker)
4344
}
4445

4546
override suspend fun truncatedListFunctions(input: TruncatedListFunctionsRequest) = TruncatedListFunctionsResponse {
46-
nextMarker = (input.marker.toIntOrZero() + 1).toString()
47-
isTruncated = input.marker.toIntOrZero() < pageCount - 1
48-
functions = generateFunctions(input.marker.toIntOrZero())
47+
val inputMarker = input.marker.toIntOrZero()
48+
49+
nextMarker = (inputMarker + 1).toString()
50+
isTruncated = inputMarker < pageCount - 1
51+
functions = generateFunctions(inputMarker)
4952
}
5053

51-
private fun generateFunctions(page: Int) = (0 until itemsPerPage).map { idx ->
52-
FunctionConfiguration { functionName = "Function page($page) item($idx)" }
54+
override suspend fun identicalTokenListFunctions(input: IdenticalTokenListFunctionsRequest) =
55+
IdenticalTokenListFunctionsResponse {
56+
val inputMarker = input.marker.toIntOrZero()
57+
58+
nextMarker = when (inputMarker) {
59+
pageCount - 1 -> input.marker // Exhausted pages, return identical input marker
60+
else -> (inputMarker + 1).toString() // Next page
61+
}
62+
63+
functions = generateFunctions(inputMarker)
64+
}
65+
66+
private fun generateFunctions(page: Int): List<FunctionConfiguration> {
67+
require(page < pageCount) { "Paginator tried to seek beyond max page $pageCount" }
68+
return (0 until itemsPerPage).map { idx ->
69+
FunctionConfiguration { functionName = "Function page($page) item($idx)" }
70+
}
5371
}
5472
}
5573

0 commit comments

Comments
 (0)