Skip to content

Commit 794b36b

Browse files
dayaffedayaffe
authored
feat: add custom trait PaginationTruncationMember (#625)
1 parent af13f00 commit 794b36b

File tree

7 files changed

+124
-17
lines changed

7 files changed

+124
-17
lines changed

‎Sources/ClientRuntime/Middleware/Steps/DeserializeStep.swift‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ public typealias DeserializeStep<OperationStackOutput> = MiddlewareStep<HttpCont
1616

1717
public let DeserializeStepId = "Deserialize"
1818

19-
public struct DeserializeStepHandler<OperationStackOutput, H: Handler>: Handler
19+
public struct DeserializeStepHandler<OperationStackOutput, H: Handler>: Handler
2020
where H.Context == HttpContext,
2121
H.Input == SdkHttpRequest,
2222
H.Output == OperationOutput<OperationStackOutput> {

‎Sources/ClientRuntime/Networking/Http/SdkHttpClient.swift‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ public class SdkHttpClient {
1212
self.engine = engine
1313
}
1414

15-
public func getHandler<OperationStackOutput>()
15+
public func getHandler<OperationStackOutput>()
1616
-> AnyHandler<SdkHttpRequest, OperationOutput<OperationStackOutput>, HttpContext> {
1717

1818
let clientHandler = ClientHandler<OperationStackOutput>(engine: engine)

‎Sources/ClientRuntime/Pagination/PaginatorSequence.swift‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,15 +12,18 @@ public struct PaginatorSequence<OperationStackInput: PaginateToken, OperationSta
1212
let input: OperationStackInput
1313
let inputKey: KeyPath<OperationStackInput, OperationStackInput.Token?>?
1414
let outputKey: KeyPath<OperationStackOutput, OperationStackInput.Token?>
15+
var isTruncatedKey: KeyPath<OperationStackOutput, Bool>?
1516
let paginationFunction: (OperationStackInput) async throws -> OperationStackOutput
1617

1718
public init(input: OperationStackInput,
1819
inputKey: KeyPath<OperationStackInput, OperationStackInput.Token?>? = nil,
1920
outputKey: KeyPath<OperationStackOutput, OperationStackInput.Token?>,
21+
isTruncatedKey: KeyPath<OperationStackOutput, Bool>? = nil,
2022
paginationFunction: @escaping (OperationStackInput) async throws -> OperationStackOutput) {
2123
self.input = input
2224
self.inputKey = inputKey
2325
self.outputKey = outputKey
26+
self.isTruncatedKey = isTruncatedKey
2427
self.paginationFunction = paginationFunction
2528
}
2629

@@ -45,6 +48,16 @@ public struct PaginatorSequence<OperationStackInput: PaginateToken, OperationSta
4548
if token != nil && token == input[keyPath: sequence.inputKey!] {
4649
break
4750
}
51+
52+
// Use isTruncatedKey from the sequence to check if pagination should continue
53+
if let isTruncatedKey = sequence.isTruncatedKey {
54+
let isTruncated = output[keyPath: isTruncatedKey]
55+
if !isTruncated {
56+
// set token to nil to break out of the next iteration
57+
token = nil
58+
}
59+
}
60+
4861
return output
4962
}
5063
return nil

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

Lines changed: 21 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -13,9 +13,11 @@ import software.amazon.smithy.model.shapes.ServiceShape
1313
import software.amazon.smithy.model.shapes.Shape
1414
import software.amazon.smithy.model.traits.PaginatedTrait
1515
import software.amazon.smithy.swift.codegen.core.CodegenContext
16+
import software.amazon.smithy.swift.codegen.customtraits.PaginationTruncationMember
1617
import software.amazon.smithy.swift.codegen.integration.ProtocolGenerator
1718
import software.amazon.smithy.swift.codegen.integration.SwiftIntegration
1819
import software.amazon.smithy.swift.codegen.model.SymbolProperty
20+
import software.amazon.smithy.swift.codegen.model.defaultName
1921
import software.amazon.smithy.swift.codegen.model.expectShape
2022
import software.amazon.smithy.swift.codegen.model.hasTrait
2123
import software.amazon.smithy.swift.codegen.model.isBoxed
@@ -55,7 +57,7 @@ class PaginatorGenerator : SwiftIntegration {
5557
service: ServiceShape,
5658
paginatedOperation: OperationShape,
5759
paginationInfo: PaginationInfo,
58-
itemDesc: ItemDescriptor?
60+
itemDesc: ItemDescriptor?,
5961
) {
6062
val serviceSymbol = ctx.symbolProvider.toSymbol(service)
6163
val outputSymbol = ctx.symbolProvider.toSymbol(paginationInfo.output)
@@ -72,7 +74,7 @@ class PaginatorGenerator : SwiftIntegration {
7274
inputSymbol,
7375
outputSymbol,
7476
paginationInfo,
75-
cursorSymbol
77+
cursorSymbol,
7678
)
7779

7880
// Optionally generate paginator when nested item is specified on the trait.
@@ -83,7 +85,7 @@ class PaginatorGenerator : SwiftIntegration {
8385
paginatedOperation,
8486
itemDesc,
8587
inputSymbol,
86-
outputSymbol
88+
outputSymbol,
8789
)
8890
}
8991
}
@@ -98,8 +100,10 @@ class PaginatorGenerator : SwiftIntegration {
98100
inputSymbol: Symbol,
99101
outputSymbol: Symbol,
100102
paginationInfo: PaginationInfo,
101-
cursorSymbol: Symbol
103+
cursorSymbol: Symbol,
102104
) {
105+
val outputShape = paginationInfo.output
106+
103107
writer.addImport(SwiftDependency.CLIENT_RUNTIME.target)
104108
val nextMarkerLiteral = paginationInfo.outputTokenMemberPath.joinToString(separator = "?.") {
105109
it.toLowerCamelCase()
@@ -122,21 +126,26 @@ class PaginatorGenerator : SwiftIntegration {
122126
this.write(docBody)
123127
}
124128
writer.openBlock(
125-
"public func \$LPaginated(input: \$N) -> \$N<\$N, \$N> {", "}",
129+
"public func \$LPaginated(input: \$N) -> \$N<\$N, \$N> {",
130+
"}",
126131
operationShape.toLowerCamelCase(),
127132
inputSymbol,
128133
ClientRuntimeTypes.Core.PaginatorSequence,
129134
inputSymbol,
130-
outputSymbol
135+
outputSymbol,
131136
) {
137+
val isTruncatedFlag = outputShape
138+
.members()
139+
.firstOrNull { it.hasTrait(PaginationTruncationMember.ID) }
140+
?.defaultName()
141+
142+
val isTruncatedPart = if (isTruncatedFlag != null) ", isTruncatedKey: \\.$isTruncatedFlag" else ""
132143
writer.write(
133-
"return \$N<\$N, \$N>(input: input, inputKey: \\\$N.$markerLiteral, outputKey: \\\$N.$nextMarkerLiteral, paginationFunction: self.\$L(input:))",
144+
"return \$N<\$N, \$N>(input: input, inputKey: \\.$markerLiteral, outputKey: \\.$nextMarkerLiteral$isTruncatedPart, paginationFunction: self.\$L(input:))",
134145
ClientRuntimeTypes.Core.PaginatorSequence,
135146
inputSymbol,
136147
outputSymbol,
137-
inputSymbol,
138-
outputSymbol,
139-
operationShape.toLowerCamelCase()
148+
operationShape.toLowerCamelCase(),
140149
)
141150
}
142151
}
@@ -212,7 +221,7 @@ private data class ItemDescriptor(
212221
val collectionLiteral: String,
213222
val itemLiteral: String,
214223
val itemPathLiteral: String,
215-
val itemSymbol: Symbol
224+
val itemSymbol: Symbol,
216225
)
217226

218227
/**
@@ -238,6 +247,6 @@ private fun getItemDescriptorOrNull(paginationInfo: PaginationInfo, ctx: Codegen
238247
collectionLiteral,
239248
itemLiteral,
240249
itemPathLiteral,
241-
ctx.symbolProvider.toSymbol(itemMember)
250+
ctx.symbolProvider.toSymbol(itemMember),
242251
)
243252
}

‎smithy-swift-codegen/src/main/kotlin/software/amazon/smithy/swift/codegen/customtraits/PaginationTruncationMember.kt‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
2+
/*
3+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
4+
* SPDX-License-Identifier: Apache-2.0
5+
*/
6+
package software.amazon.smithy.swift.codegen.customtraits
7+
8+
import software.amazon.smithy.model.node.Node
9+
import software.amazon.smithy.model.node.ObjectNode
10+
import software.amazon.smithy.model.shapes.ShapeId
11+
import software.amazon.smithy.model.traits.AnnotationTrait
12+
13+
/**
14+
* Indicates the annotated member is a truncation indicator which conveys a non-standard termination condition for
15+
* pagination.
16+
*/
17+
class PaginationTruncationMember(node: ObjectNode) : AnnotationTrait(ID, node) {
18+
companion object {
19+
val ID: ShapeId = ShapeId.from("software.amazon.smithy.swift.codegen.synthetic#paginationTruncationMember")
20+
}
21+
22+
constructor() : this(Node.objectNode())
23+
}

‎smithy-swift-codegen/src/test/kotlin/PaginatorGeneratorTest.kt‎

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ class PaginatorGeneratorTest {
2525
/// - input: A `[ListFunctionsInput]` to start pagination
2626
/// - Returns: An `AsyncSequence` that can iterate over `ListFunctionsOutput`
2727
public func listFunctionsPaginated(input: ListFunctionsInput) -> ClientRuntime.PaginatorSequence<ListFunctionsInput, ListFunctionsOutput> {
28-
return ClientRuntime.PaginatorSequence<ListFunctionsInput, ListFunctionsOutput>(input: input, inputKey: \ListFunctionsInput.marker, outputKey: \ListFunctionsOutput.nextMarker, paginationFunction: self.listFunctions(input:))
28+
return ClientRuntime.PaginatorSequence<ListFunctionsInput, ListFunctionsOutput>(input: input, inputKey: \.marker, outputKey: \.nextMarker, paginationFunction: self.listFunctions(input:))
2929
}
3030
}
3131
@@ -58,7 +58,7 @@ class PaginatorGeneratorTest {
5858
/// - input: A `[ListFunctionsInput]` to start pagination
5959
/// - Returns: An `AsyncSequence` that can iterate over `ListFunctionsOutput`
6060
public func listFunctionsPaginated(input: ListFunctionsInput) -> ClientRuntime.PaginatorSequence<ListFunctionsInput, ListFunctionsOutput> {
61-
return ClientRuntime.PaginatorSequence<ListFunctionsInput, ListFunctionsOutput>(input: input, inputKey: \ListFunctionsInput.marker, outputKey: \ListFunctionsOutput.nextMarker, paginationFunction: self.listFunctions(input:))
61+
return ClientRuntime.PaginatorSequence<ListFunctionsInput, ListFunctionsOutput>(input: input, inputKey: \.marker, outputKey: \.nextMarker, paginationFunction: self.listFunctions(input:))
6262
}
6363
}
6464
@@ -100,7 +100,7 @@ class PaginatorGeneratorTest {
100100
/// - input: A `[PaginatedMapInput]` to start pagination
101101
/// - Returns: An `AsyncSequence` that can iterate over `PaginatedMapOutput`
102102
public func paginatedMapPaginated(input: PaginatedMapInput) -> ClientRuntime.PaginatorSequence<PaginatedMapInput, PaginatedMapOutput> {
103-
return ClientRuntime.PaginatorSequence<PaginatedMapInput, PaginatedMapOutput>(input: input, inputKey: \PaginatedMapInput.nextToken, outputKey: \PaginatedMapOutput.inner?.token, paginationFunction: self.paginatedMap(input:))
103+
return ClientRuntime.PaginatorSequence<PaginatedMapInput, PaginatedMapOutput>(input: input, inputKey: \.nextToken, outputKey: \.inner?.token, paginationFunction: self.paginatedMap(input:))
104104
}
105105
}
106106
@@ -125,6 +125,18 @@ class PaginatorGeneratorTest {
125125
contents.shouldContainOnlyOnce(expectedCode)
126126
}
127127

128+
@Test
129+
fun testRenderPaginatorTruncatable() {
130+
val context = setupTests("pagination-truncation.smithy", "software.amazon.smithy.swift.codegen.synthetic#Lambda")
131+
val contents = getFileContents(context.manifest, "/Test/Paginators.swift")
132+
val expected = """
133+
public func listFunctionsTruncatedPaginated(input: ListFunctionsTruncatedInput) -> ClientRuntime.PaginatorSequence<ListFunctionsTruncatedInput, ListFunctionsTruncatedOutput> {
134+
return ClientRuntime.PaginatorSequence<ListFunctionsTruncatedInput, ListFunctionsTruncatedOutput>(input: input, inputKey: \.marker, outputKey: \.nextMarker, isTruncatedKey: \.isTruncated, paginationFunction: self.listFunctionsTruncated(input:))
135+
}
136+
"""
137+
contents.shouldContainOnlyOnce(expected)
138+
}
139+
128140
private fun setupTests(smithyFile: String, serviceShapeId: String): TestContext {
129141
val context = TestContext.initContextFrom(smithyFile, serviceShapeId, MockHttpRestJsonProtocolGenerator()) { model ->
130142
model.defaultSettings(serviceShapeId, "Test", "2019-12-16", "Test")

‎smithy-swift-codegen/src/test/resources/pagination-truncation.smithy‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
$version: "1.0"
2+
3+
namespace software.amazon.smithy.swift.codegen.synthetic
4+
5+
use aws.protocols#restJson1
6+
7+
@trait(selector: "*")
8+
structure paginationTruncationMember { }
9+
10+
service Lambda {
11+
operations: [ListFunctionsTruncated]
12+
}
13+
14+
list FunctionConfigurationList {
15+
member: FunctionConfiguration
16+
}
17+
18+
structure FunctionConfiguration {
19+
functionName: String
20+
}
21+
22+
@paginated(
23+
inputToken: "marker",
24+
outputToken: "nextMarker",
25+
pageSize: "maxItems"
26+
)
27+
@readonly
28+
@http(method: "GET", uri: "/functions/truncated", code: 200)
29+
operation ListFunctionsTruncated {
30+
input: ListFunctionsRequestTruncated,
31+
output: ListFunctionsResponseTruncated
32+
}
33+
34+
structure ListFunctionsRequestTruncated {
35+
@httpQuery("FunctionVersion")
36+
functionVersion: String,
37+
@httpQuery("Marker")
38+
marker: String,
39+
@httpQuery("MasterRegion")
40+
masterRegion: String,
41+
@httpQuery("MaxItems")
42+
maxItems: Integer,
43+
}
44+
45+
structure ListFunctionsResponseTruncated {
46+
Functions: FunctionConfigurationList,
47+
@paginationTruncationMember
48+
IsTruncated: Boolean,
49+
nextMarker: String
50+
}

0 commit comments

Comments
 (0)