Skip to content

Commit 0ac9bfb

Browse files
kneekey23kneekey23
authored
feat: pagination (#410)
1 parent f103928 commit 0ac9bfb

File tree

12 files changed

+576
-7
lines changed

12 files changed

+576
-7
lines changed

‎Packages/ClientRuntime/Sources/Pagination/PaginateToken.swift‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
//
2+
// Copyright Amazon.com Inc. or its affiliates.
3+
// All Rights Reserved.
4+
//
5+
// SPDX-License-Identifier: Apache-2.0
6+
//
7+
8+
/// protocol for all Inputs that can be paginated.
9+
/// Adds an initializer that does a copy but inserts a new integer based pagination token
10+
public protocol PaginateToken {
11+
associatedtype Token
12+
func usingPaginationToken(_ token: Token) -> Self
13+
}

‎Packages/ClientRuntime/Sources/Pagination/PaginatorSequence.swift‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
//
2+
// Copyright Amazon.com Inc. or its affiliates.
3+
// All Rights Reserved.
4+
//
5+
// SPDX-License-Identifier: Apache-2.0
6+
//
7+
8+
public struct PaginatorSequence<Input: PaginateToken,
9+
Output: HttpResponseBinding>: AsyncSequence
10+
where Input.Token: Equatable {
11+
public typealias Element = Output
12+
let input: Input
13+
let inputKey: KeyPath<Input, Input.Token?>?
14+
let outputKey: KeyPath<Output, Input.Token?>
15+
let paginationFunction: (Input) async throws -> Output
16+
17+
public init(input: Input,
18+
inputKey: KeyPath<Input, Input.Token?>? = nil,
19+
outputKey: KeyPath<Output, Input.Token?>,
20+
paginationFunction: @escaping (Input) async throws -> Output) {
21+
self.input = input
22+
self.inputKey = inputKey
23+
self.outputKey = outputKey
24+
self.paginationFunction = paginationFunction
25+
}
26+
27+
public struct PaginationIterator: AsyncIteratorProtocol {
28+
var input: Input
29+
let sequence: PaginatorSequence
30+
var token: Input.Token?
31+
var isFirstPage: Bool = true
32+
33+
// swiftlint:disable force_cast
34+
public mutating func next() async throws -> Output? {
35+
while token != nil || isFirstPage {
36+
37+
if let token = token,
38+
(token is String && !(token as! String).isEmpty) ||
39+
(token is [String: Any] && !(token as! [String: Any]).isEmpty) {
40+
self.input = input.usingPaginationToken(token)
41+
}
42+
let output = try await sequence.paginationFunction(input)
43+
isFirstPage = false
44+
token = output[keyPath: sequence.outputKey]
45+
if token == input[keyPath: sequence.inputKey!] {
46+
break
47+
}
48+
return output
49+
}
50+
return nil
51+
}
52+
}
53+
54+
public func makeAsyncIterator() -> PaginationIterator {
55+
PaginationIterator(input: input, sequence: self)
56+
}
57+
}

‎Packages/ClientRuntime/Sources/PrimitiveTypeExtensions/AsyncSequence+Extension.swift‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
//
2+
// Copyright Amazon.com Inc. or its affiliates.
3+
// All Rights Reserved.
4+
//
5+
// SPDX-License-Identifier: Apache-2.0
6+
//
7+
8+
extension AsyncSequence {
9+
public func asyncCompactMap<T>(
10+
_ transform: (Element) -> [T]?
11+
) async rethrows -> [T] {
12+
var values = [T]()
13+
14+
for try await element in self {
15+
if let element = transform(element) {
16+
values.append(contentsOf: element)
17+
}
18+
}
19+
20+
return values
21+
}
22+
}

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,8 @@ object ClientRuntimeTypes {
9393
val SDKRuntimeConfiguration = runtimeSymbol("SDKRuntimeConfiguration")
9494
val DefaultSDKRuntimeConfiguration = runtimeSymbol("DefaultSDKRuntimeConfiguration")
9595
val DateFormatter = runtimeSymbol("DateFormatter")
96+
val PaginateToken = runtimeSymbol("PaginateToken")
97+
val PaginatorSequence = runtimeSymbol("PaginatorSequence")
9698
}
9799
}
98100

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

Lines changed: 220 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,220 @@
1+
package software.amazon.smithy.swift.codegen
2+
3+
import software.amazon.smithy.codegen.core.CodegenException
4+
import software.amazon.smithy.codegen.core.Symbol
5+
import software.amazon.smithy.codegen.core.SymbolProvider
6+
import software.amazon.smithy.model.Model
7+
import software.amazon.smithy.model.knowledge.PaginatedIndex
8+
import software.amazon.smithy.model.knowledge.PaginationInfo
9+
import software.amazon.smithy.model.shapes.OperationShape
10+
import software.amazon.smithy.model.shapes.ServiceShape
11+
import software.amazon.smithy.model.traits.PaginatedTrait
12+
import software.amazon.smithy.swift.codegen.core.CodegenContext
13+
import software.amazon.smithy.swift.codegen.integration.SwiftIntegration
14+
import software.amazon.smithy.swift.codegen.model.camelCaseName
15+
import software.amazon.smithy.swift.codegen.model.expectShape
16+
import software.amazon.smithy.swift.codegen.model.hasTrait
17+
import software.amazon.smithy.swift.codegen.utils.toCamelCase
18+
19+
/**
20+
* Generate paginators for supporting operations. See
21+
* https://awslabs.github.io/smithy/1.0/spec/core/behavior-traits.html#paginated-trait for details.
22+
*/
23+
class PaginatorGenerator : SwiftIntegration {
24+
override fun enabledForService(model: Model, settings: SwiftSettings): Boolean =
25+
model.operationShapes.any { it.hasTrait<PaginatedTrait>() }
26+
27+
override fun writeAdditionalFiles(ctx: CodegenContext, delegator: SwiftDelegator) {
28+
val service = ctx.model.expectShape<ServiceShape>(ctx.settings.service)
29+
val paginatedIndex = PaginatedIndex.of(ctx.model)
30+
31+
delegator.useFileWriter("${ctx.settings.moduleName}/Paginators.swift") { writer ->
32+
val paginatedOperations = service.allOperations
33+
.map { ctx.model.expectShape<OperationShape>(it) }
34+
.filter { operationShape -> operationShape.hasTrait(PaginatedTrait.ID) }
35+
36+
paginatedOperations.forEach { paginatedOperation ->
37+
val paginationInfo = paginatedIndex.getPaginationInfo(service, paginatedOperation).getOrNull()
38+
?: throw CodegenException("Unexpectedly unable to get PaginationInfo from $service $paginatedOperation")
39+
val paginationItemInfo = getItemDescriptorOrNull(paginationInfo, ctx)
40+
renderPaginatorForOperation(writer, ctx, service, paginatedOperation, paginationInfo, paginationItemInfo)
41+
}
42+
}
43+
}
44+
45+
// Render paginator(s) for operation
46+
private fun renderPaginatorForOperation(
47+
writer: SwiftWriter,
48+
ctx: CodegenContext,
49+
service: ServiceShape,
50+
paginatedOperation: OperationShape,
51+
paginationInfo: PaginationInfo,
52+
itemDesc: ItemDescriptor?
53+
) {
54+
val serviceSymbol = ctx.symbolProvider.toSymbol(service)
55+
val outputSymbol = ctx.symbolProvider.toSymbol(paginationInfo.output)
56+
val inputSymbol = ctx.symbolProvider.toSymbol(paginationInfo.input)
57+
val cursorMember = ctx.model.getShape(paginationInfo.inputTokenMember.target).get()
58+
val cursorSymbol = ctx.symbolProvider.toSymbol(cursorMember)
59+
60+
renderResponsePaginator(
61+
writer,
62+
ctx.model,
63+
ctx.symbolProvider,
64+
serviceSymbol,
65+
paginatedOperation,
66+
inputSymbol,
67+
outputSymbol,
68+
paginationInfo,
69+
cursorSymbol
70+
)
71+
72+
// Optionally generate paginator when nested item is specified on the trait.
73+
if (itemDesc != null) {
74+
renderItemPaginator(
75+
writer,
76+
service,
77+
paginatedOperation,
78+
itemDesc,
79+
inputSymbol,
80+
outputSymbol
81+
)
82+
}
83+
}
84+
85+
// Generate the paginator that iterates over responses
86+
private fun renderResponsePaginator(
87+
writer: SwiftWriter,
88+
model: Model,
89+
symbolProvider: SymbolProvider,
90+
serviceSymbol: Symbol,
91+
operationShape: OperationShape,
92+
inputSymbol: Symbol,
93+
outputSymbol: Symbol,
94+
paginationInfo: PaginationInfo,
95+
cursorSymbol: Symbol
96+
) {
97+
writer.addImport(SwiftDependency.CLIENT_RUNTIME.target)
98+
val nextMarkerLiteral = paginationInfo.outputTokenMemberPath.joinToString(separator = "?.") {
99+
it.camelCaseName()
100+
}
101+
val markerLiteral = paginationInfo.inputTokenMember.camelCaseName()
102+
val markerLiteralShape = model.expectShape(paginationInfo.inputTokenMember.target)
103+
val markerLiteralSymbol = symbolProvider.toSymbol(markerLiteralShape)
104+
val docBody = """
105+
Paginate over `[${outputSymbol.name}]` results.
106+
107+
When this operation is called, an `AsyncSequence` is created. AsyncSequences are lazy so no service
108+
calls are made until the sequence is iterated over. This also means there is no guarantee that the request is valid
109+
until then. If there are errors in your request, you will see the failures only after you start iterating.
110+
- Parameters:
111+
- input: A `[${inputSymbol.name}]` to start pagination
112+
- Returns: An `AsyncSequence` that can iterate over `${outputSymbol.name}`
113+
""".trimIndent()
114+
writer.write("")
115+
writer.writeSingleLineDocs {
116+
this.write(docBody)
117+
}
118+
119+
writer.openBlock("extension \$L {", "}", serviceSymbol.name) {
120+
writer.openBlock(
121+
"public func \$LPaginated(input: \$N) -> \$N<\$N, \$N> {", "}",
122+
operationShape.camelCaseName(),
123+
inputSymbol,
124+
ClientRuntimeTypes.Core.PaginatorSequence,
125+
inputSymbol,
126+
outputSymbol
127+
) {
128+
writer.write(
129+
"return \$N<\$N, \$N>(input: input, inputKey: \\\$N.$markerLiteral, outputKey: \\\$N.$nextMarkerLiteral, paginationFunction: self.\$L(input:))",
130+
ClientRuntimeTypes.Core.PaginatorSequence,
131+
inputSymbol,
132+
outputSymbol,
133+
inputSymbol,
134+
outputSymbol,
135+
operationShape.camelCaseName()
136+
)
137+
}
138+
}
139+
140+
writer.write("")
141+
142+
writer.openBlock("extension \$N: \$N {", "}", inputSymbol, ClientRuntimeTypes.Core.PaginateToken) {
143+
writer.openBlock("public func usingPaginationToken(_ token: \$N) -> \$N {", "}", markerLiteralSymbol, inputSymbol) {
144+
writer.writeInline("return ")
145+
.call {
146+
val inputShape = model.expectShape(operationShape.input.get())
147+
writer.writeInline("\$N(", inputSymbol)
148+
.indent()
149+
.call {
150+
val sortedMembers = inputShape.members().sortedBy { it.camelCaseName() }
151+
for ((index, member) in sortedMembers.withIndex()) {
152+
if (member.memberName.toCamelCase() != markerLiteral) {
153+
writer.writeInline("\n\$L: \$L", member.camelCaseName(), "self.${member.camelCaseName()}")
154+
} else {
155+
writer.writeInline("\n\$L: \$L", member.camelCaseName(), "token")
156+
}
157+
if (index < sortedMembers.size - 1) {
158+
writer.writeInline(",")
159+
}
160+
}
161+
}
162+
.dedent()
163+
.writeInline("\n)")
164+
}
165+
}
166+
}
167+
}
168+
169+
// Generate a paginator that iterates over the model-specified item
170+
private fun renderItemPaginator(
171+
writer: SwiftWriter,
172+
serviceShape: ServiceShape,
173+
operationShape: OperationShape,
174+
itemDesc: ItemDescriptor,
175+
inputSymbol: Symbol,
176+
outputSymbol: Symbol,
177+
) {
178+
writer.write("")
179+
val docBody = """
180+
This paginator transforms the `AsyncSequence` returned by `${operationShape.camelCaseName()}Paginated`
181+
to access the nested member `${itemDesc.itemSymbol.fullName}`
182+
- Returns: `${itemDesc.itemSymbol.fullName}`
183+
""".trimIndent()
184+
185+
writer.writeSingleLineDocs {
186+
this.write(docBody)
187+
}
188+
189+
writer.openBlock("extension PaginatorSequence where Input == \$N, Output == \$N {", "}", inputSymbol, outputSymbol) {
190+
writer.openBlock("func \$L() async throws -> \$N {", "}", itemDesc.itemLiteral, itemDesc.itemSymbol) {
191+
writer.write("return try await self.asyncCompactMap { item in item.\$L }", itemDesc.itemPathLiteral)
192+
}
193+
}
194+
}
195+
}
196+
197+
/**
198+
* Model info necessary to codegen paginator item
199+
*/
200+
private data class ItemDescriptor(
201+
val itemLiteral: String,
202+
val itemPathLiteral: String,
203+
val itemSymbol: Symbol
204+
)
205+
206+
/**
207+
* Return an [ItemDescriptor] if model supplies, otherwise null
208+
*/
209+
private fun getItemDescriptorOrNull(paginationInfo: PaginationInfo, ctx: CodegenContext): ItemDescriptor? {
210+
val itemMemberId = paginationInfo.itemsMemberPath?.lastOrNull()?.target ?: return null
211+
val itemLiteral = paginationInfo.itemsMemberPath!!.last()!!.camelCaseName()
212+
val itemPathLiteral = paginationInfo.itemsMemberPath.joinToString(separator = "?.") { it.camelCaseName() }
213+
val itemMember = ctx.model.expectShape(itemMemberId)
214+
215+
return ItemDescriptor(
216+
itemLiteral,
217+
itemPathLiteral,
218+
ctx.symbolProvider.toSymbol(itemMember)
219+
)
220+
}

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import software.amazon.smithy.swift.codegen.customtraits.NestedTrait
1919
import software.amazon.smithy.swift.codegen.customtraits.SwiftBoxTrait
2020
import software.amazon.smithy.swift.codegen.integration.ProtocolGenerator
2121
import software.amazon.smithy.swift.codegen.integration.SectionId
22+
import software.amazon.smithy.swift.codegen.model.camelCaseName
2223
import software.amazon.smithy.swift.codegen.model.expectShape
2324
import software.amazon.smithy.swift.codegen.model.getTrait
2425
import software.amazon.smithy.swift.codegen.model.hasTrait
@@ -35,7 +36,7 @@ class StructureGenerator(
3536
private val serviceErrorProtocolSymbol: Symbol? = null
3637
) {
3738

38-
private val membersSortedByName: List<MemberShape> = shape.allMembers.values.sortedBy { symbolProvider.toMemberName(it) }
39+
private val membersSortedByName: List<MemberShape> = shape.allMembers.values.sortedBy { it.camelCaseName() }
3940
private var memberShapeDataContainer: MutableMap<MemberShape, Pair<String, Symbol>> = mutableMapOf()
4041

4142
init {

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -75,14 +75,14 @@ fun Shape.capitalizedName(): String {
7575
return StringUtils.capitalize(this.id.name)
7676
}
7777

78-
fun Shape.defaultName(serviceShape: ServiceShape?): String {
78+
fun Shape.defaultName(serviceShape: ServiceShape? = null): String {
7979
return serviceShape?.let {
8080
StringUtils.capitalize(id.getName(it))
8181
} ?: run {
8282
StringUtils.capitalize(this.id.name)
8383
}
8484
}
85-
85+
fun MemberShape.camelCaseName(): String = StringUtils.uncapitalize(this.memberName)
8686
fun Shape.camelCaseName(): String = StringUtils.uncapitalize(this.id.name)
8787

8888
fun MemberShape.defaultValue(symbolProvider: SymbolProvider): String? {

‎smithy-swift-codegen/src/main/resources/META-INF/services/software.amazon.smithy.swift.codegen.integration.SwiftIntegration‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
software.amazon.smithy.swift.codegen.PaginatorGenerator

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -70,16 +70,16 @@ class HashableShapeTransformerTests {
7070
Assertions.assertNotNull(hashableShapeInput)
7171
val expected = """
7272
public struct HashableShapesInput: Swift.Equatable {
73-
public var `set`: Swift.Set<ExampleClientTypes.HashableStructure>?
7473
public var bar: Swift.String?
74+
public var `set`: Swift.Set<ExampleClientTypes.HashableStructure>?
7575
7676
public init (
77-
`set`: Swift.Set<ExampleClientTypes.HashableStructure>? = nil,
78-
bar: Swift.String? = nil
77+
bar: Swift.String? = nil,
78+
`set`: Swift.Set<ExampleClientTypes.HashableStructure>? = nil
7979
)
8080
{
81-
self.`set` = `set`
8281
self.bar = bar
82+
self.`set` = `set`
8383
}
8484
}
8585
""".trimIndent()

0 commit comments

Comments
 (0)