Skip to content

Commit 51480df

Browse files
trivikrtrivikr
authored
feat(lib-dynamodb): add pagination (#3069)
1 parent 31f3478 commit 51480df

File tree

8 files changed

+375
-9
lines changed

8 files changed

+375
-9
lines changed

‎codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AddDocumentClientPlugin.java

Lines changed: 37 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import software.amazon.smithy.model.shapes.OperationShape;
2929
import software.amazon.smithy.model.shapes.ServiceShape;
3030
import software.amazon.smithy.model.shapes.Shape;
31+
import software.amazon.smithy.model.traits.PaginatedTrait;
3132
import software.amazon.smithy.typescript.codegen.TypeScriptSettings;
3233
import software.amazon.smithy.typescript.codegen.TypeScriptWriter;
3334
import software.amazon.smithy.typescript.codegen.integration.TypeScriptIntegration;
@@ -49,14 +50,13 @@ public void writeAdditionalFiles(
4950
) {
5051
ServiceShape service = settings.getService(model);
5152
if (testServiceId(service, "DynamoDB")) {
52-
String docClientPrefix = "doc-client-";
5353
Set<OperationShape> containedOperations =
5454
new TreeSet<>(TopDownIndex.of(model).getContainedOperations(service));
5555
List<OperationShape> overridenOperationsList = new ArrayList<>();
5656

5757
for (OperationShape operation : containedOperations) {
5858
String operationName = symbolProvider.toSymbol(operation).getName();
59-
String commandFileName = String.format("%s%s/%s.ts", docClientPrefix,
59+
String commandFileName = String.format("%s%s/%s.ts", DocumentClientUtils.DOC_CLIENT_PREFIX,
6060
DocumentClientUtils.CLIENT_COMMANDS_FOLDER, DocumentClientUtils.getModifiedName(operationName));
6161

6262
if (DocumentClientUtils.containsAttributeValue(model, symbolProvider, operation)) {
@@ -65,16 +65,25 @@ public void writeAdditionalFiles(
6565
writer -> new DocumentClientCommandGenerator(
6666
settings, model, operation, symbolProvider, writer).run()
6767
);
68+
69+
if (operation.hasTrait(PaginatedTrait.ID)) {
70+
String paginationFileName = DocumentClientPaginationGenerator.getOutputFilelocation(operation);
71+
writerFactory.accept(paginationFileName, paginationWriter ->
72+
new DocumentClientPaginationGenerator(model, service, operation, symbolProvider,
73+
paginationWriter).run());
74+
}
6875
}
6976
}
7077

71-
writerFactory.accept(String.format("%s%s.ts", docClientPrefix, DocumentClientUtils.CLIENT_NAME),
78+
writerFactory.accept(String.format("%s%s.ts", DocumentClientUtils.DOC_CLIENT_PREFIX,
79+
DocumentClientUtils.CLIENT_NAME),
7280
writer -> new DocumentBareBonesClientGenerator(settings, model, symbolProvider, writer).run());
7381

74-
writerFactory.accept(String.format("%s%s.ts", docClientPrefix, DocumentClientUtils.CLIENT_FULL_NAME),
82+
writerFactory.accept(String.format("%s%s.ts", DocumentClientUtils.DOC_CLIENT_PREFIX,
83+
DocumentClientUtils.CLIENT_FULL_NAME),
7584
writer -> new DocumentAggregatedClientGenerator(settings, model, symbolProvider, writer).run());
7685

77-
writerFactory.accept(String.format("%s%s/index.ts", docClientPrefix,
86+
writerFactory.accept(String.format("%s%s/index.ts", DocumentClientUtils.DOC_CLIENT_PREFIX,
7887
DocumentClientUtils.CLIENT_COMMANDS_FOLDER), writer -> {
7988
for (OperationShape operation : overridenOperationsList) {
8089
String operationFileName = DocumentClientUtils.getModifiedName(
@@ -83,19 +92,38 @@ public void writeAdditionalFiles(
8392
writer.write("export * from './$L';", operationFileName);
8493
}
8594
});
86-
writerFactory.accept(String.format("%sindex.ts", docClientPrefix), writer -> {
95+
96+
writerFactory.accept(String.format("%s%s/index.ts", DocumentClientUtils.DOC_CLIENT_PREFIX,
97+
DocumentClientPaginationGenerator.PAGINATION_FOLDER), writer -> {
98+
writer.write("export * from './Interfaces';");
99+
for (OperationShape operation : overridenOperationsList) {
100+
if (operation.hasTrait(PaginatedTrait.ID)) {
101+
String paginationFileName =
102+
DocumentClientUtils.getModifiedName(operation.getId().getName()) + "Paginator";
103+
writer.write("export * from './$L';", paginationFileName);
104+
}
105+
}
106+
});
107+
108+
String paginationInterfaceFileName = DocumentClientPaginationGenerator.getInterfaceFilelocation();
109+
writerFactory.accept(paginationInterfaceFileName, paginationWriter ->
110+
DocumentClientPaginationGenerator.generateServicePaginationInterfaces(paginationWriter));
111+
112+
writerFactory.accept(String.format("%sindex.ts", DocumentClientUtils.DOC_CLIENT_PREFIX), writer -> {
87113
writer.write("export * from './commands';");
114+
writer.write("export * from './pagination';");
88115
writer.write("export * from './$L';", DocumentClientUtils.CLIENT_NAME);
89116
writer.write("export * from './$L';", DocumentClientUtils.CLIENT_FULL_NAME);
90117
});
91118

92-
String utilsFileLocation = String.format("%s%s", docClientPrefix, DocumentClientUtils.CLIENT_UTILS_FILE);
93-
writerFactory.accept(String.format("%s%s/%s.ts", docClientPrefix,
119+
String utilsFileLocation = String.format("%s%s", DocumentClientUtils.DOC_CLIENT_PREFIX,
120+
DocumentClientUtils.CLIENT_UTILS_FILE);
121+
writerFactory.accept(String.format("%s%s/%s.ts", DocumentClientUtils.DOC_CLIENT_PREFIX,
94122
DocumentClientUtils.CLIENT_COMMANDS_FOLDER, DocumentClientUtils.CLIENT_UTILS_FILE), writer -> {
95123
writer.write(IoUtils.readUtf8Resource(AddDocumentClientPlugin.class,
96124
String.format("%s.ts", utilsFileLocation)));
97125
});
98-
writerFactory.accept(String.format("%s%s/%s.spec.ts", docClientPrefix,
126+
writerFactory.accept(String.format("%s%s/%s.spec.ts", DocumentClientUtils.DOC_CLIENT_PREFIX,
99127
DocumentClientUtils.CLIENT_COMMANDS_FOLDER, DocumentClientUtils.CLIENT_UTILS_FILE), writer -> {
100128
writer.write(IoUtils.readUtf8Resource(AddDocumentClientPlugin.class,
101129
String.format("%s.spec.ts", utilsFileLocation)));

‎codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/DocumentClientPaginationGenerator.java

Lines changed: 215 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,215 @@
1+
/*
2+
* Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License").
5+
* You may not use this file except in compliance with the License.
6+
* A copy of the License is located at
7+
*
8+
* http://aws.amazon.com/apache2.0
9+
*
10+
* or in the "license" file accompanying this file. This file is distributed
11+
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12+
* express or implied. See the License for the specific language governing
13+
* permissions and limitations under the License.
14+
*/
15+
16+
package software.amazon.smithy.aws.typescript.codegen;
17+
18+
import java.nio.file.Paths;
19+
import java.util.Optional;
20+
import software.amazon.smithy.codegen.core.CodegenException;
21+
import software.amazon.smithy.codegen.core.Symbol;
22+
import software.amazon.smithy.codegen.core.SymbolProvider;
23+
import software.amazon.smithy.model.Model;
24+
import software.amazon.smithy.model.knowledge.PaginatedIndex;
25+
import software.amazon.smithy.model.knowledge.PaginationInfo;
26+
import software.amazon.smithy.model.shapes.OperationShape;
27+
import software.amazon.smithy.model.shapes.ServiceShape;
28+
import software.amazon.smithy.typescript.codegen.TypeScriptWriter;
29+
import software.amazon.smithy.utils.SmithyInternalApi;
30+
31+
@SmithyInternalApi
32+
final class DocumentClientPaginationGenerator implements Runnable {
33+
34+
static final String PAGINATION_FOLDER = "pagination";
35+
36+
private final TypeScriptWriter writer;
37+
private final PaginationInfo paginatedInfo;
38+
39+
private final String operationTypeName;
40+
private final String inputTypeName;
41+
private final String outputTypeName;
42+
43+
private final String operationName;
44+
private final String methodName;
45+
private final String paginationType;
46+
47+
DocumentClientPaginationGenerator(
48+
Model model,
49+
ServiceShape service,
50+
OperationShape operation,
51+
SymbolProvider symbolProvider,
52+
TypeScriptWriter writer
53+
) {
54+
this.writer = writer;
55+
56+
Symbol operationSymbol = symbolProvider.toSymbol(operation);
57+
Symbol inputSymbol = symbolProvider.toSymbol(operation).expectProperty("inputType", Symbol.class);
58+
Symbol outputSymbol = symbolProvider.toSymbol(operation).expectProperty("outputType", Symbol.class);
59+
60+
this.operationTypeName = DocumentClientUtils.getModifiedName(operationSymbol.getName());
61+
this.inputTypeName = DocumentClientUtils.getModifiedName(inputSymbol.getName());
62+
this.outputTypeName = DocumentClientUtils.getModifiedName(outputSymbol.getName());
63+
64+
// e.g. listObjects
65+
this.operationName = operationTypeName.replace("Command", "");
66+
this.methodName = Character.toLowerCase(operationName.charAt(0)) + operationName.substring(1);
67+
this.paginationType = DocumentClientUtils.CLIENT_FULL_NAME + "PaginationConfiguration";
68+
69+
PaginatedIndex paginatedIndex = PaginatedIndex.of(model);
70+
Optional<PaginationInfo> paginationInfo = paginatedIndex.getPaginationInfo(service, operation);
71+
this.paginatedInfo = paginationInfo.orElseThrow(() -> {
72+
return new CodegenException("Expected Paginator to have pagination information.");
73+
});
74+
}
75+
76+
@Override
77+
public void run() {
78+
// Import Service Types
79+
String commandFileLocation = Paths.get(".", DocumentClientUtils.CLIENT_COMMANDS_FOLDER,
80+
DocumentClientUtils.getModifiedName(operationTypeName)).toString();
81+
writer.addImport(operationTypeName, operationTypeName, commandFileLocation);
82+
writer.addImport(inputTypeName, inputTypeName, commandFileLocation);
83+
writer.addImport(outputTypeName, outputTypeName, commandFileLocation);
84+
writer.addImport(
85+
DocumentClientUtils.CLIENT_NAME,
86+
DocumentClientUtils.CLIENT_NAME,
87+
Paths.get(".", DocumentClientUtils.CLIENT_NAME).toString());
88+
writer.addImport(
89+
DocumentClientUtils.CLIENT_FULL_NAME,
90+
DocumentClientUtils.CLIENT_FULL_NAME,
91+
Paths.get(".", DocumentClientUtils.CLIENT_FULL_NAME).toString());
92+
93+
// Import Pagination types
94+
writer.addImport("Paginator", "Paginator", "@aws-sdk/types");
95+
writer.addImport(paginationType, paginationType,
96+
Paths.get(".", getInterfaceFilelocation().replace(".ts", "")).toString());
97+
98+
writeCommandRequest();
99+
writeMethodRequest();
100+
writePager();
101+
}
102+
103+
static String getOutputFilelocation(OperationShape operation) {
104+
return String.format("%s%s/%s.ts", DocumentClientUtils.DOC_CLIENT_PREFIX,
105+
DocumentClientPaginationGenerator.PAGINATION_FOLDER,
106+
DocumentClientUtils.getModifiedName(operation.getId().getName()) + "Paginator");
107+
}
108+
109+
static String getInterfaceFilelocation() {
110+
return String.format("%s%s/%s.ts", DocumentClientUtils.DOC_CLIENT_PREFIX,
111+
DocumentClientPaginationGenerator.PAGINATION_FOLDER, "Interfaces");
112+
}
113+
114+
static void generateServicePaginationInterfaces(TypeScriptWriter writer) {
115+
writer.addImport("PaginationConfiguration", "PaginationConfiguration", "@aws-sdk/types");
116+
117+
writer.addImport(
118+
DocumentClientUtils.CLIENT_NAME,
119+
DocumentClientUtils.CLIENT_NAME,
120+
Paths.get(".", DocumentClientUtils.CLIENT_NAME).toString());
121+
writer.addImport(
122+
DocumentClientUtils.CLIENT_FULL_NAME,
123+
DocumentClientUtils.CLIENT_FULL_NAME,
124+
Paths.get(".", DocumentClientUtils.CLIENT_FULL_NAME).toString());
125+
126+
writer.openBlock("export interface $LPaginationConfiguration extends PaginationConfiguration {",
127+
"}", DocumentClientUtils.CLIENT_FULL_NAME, () -> {
128+
writer.write("client: $L | $L;", DocumentClientUtils.CLIENT_FULL_NAME, DocumentClientUtils.CLIENT_NAME);
129+
});
130+
}
131+
132+
private String destructurePath(String path) {
133+
return "." + path.replace(".", "!.");
134+
}
135+
136+
private void writePager() {
137+
String inputTokenName = paginatedInfo.getPaginatedTrait().getInputToken().get();
138+
String outputTokenName = paginatedInfo.getPaginatedTrait().getOutputToken().get();
139+
140+
writer.openBlock(
141+
"export async function* paginate$L(config: $L, input: $L, ...additionalArguments: any): Paginator<$L>{",
142+
"}", operationName, paginationType, inputTypeName, outputTypeName, () -> {
143+
String destructuredInputTokenName = destructurePath(inputTokenName);
144+
writer.write("// ToDo: replace with actual type instead of typeof input$L", destructuredInputTokenName);
145+
writer.write("let token: typeof input$L | undefined = config.startingToken || undefined;",
146+
destructuredInputTokenName);
147+
148+
writer.write("let hasNext = true;");
149+
writer.write("let page: $L;", outputTypeName);
150+
writer.openBlock("while (hasNext) {", "}", () -> {
151+
writer.write("input$L = token;", destructuredInputTokenName);
152+
153+
if (paginatedInfo.getPageSizeMember().isPresent()) {
154+
String pageSize = paginatedInfo.getPageSizeMember().get().getMemberName();
155+
writer.write("input[$S] = config.pageSize;", pageSize);
156+
}
157+
158+
writer.openBlock("if (config.client instanceof $L) {", "}", DocumentClientUtils.CLIENT_FULL_NAME,
159+
() -> {
160+
writer.write("page = await makePagedRequest(config.client, input, ...additionalArguments);");
161+
}
162+
);
163+
writer.openBlock("else if (config.client instanceof $L) {", "}", DocumentClientUtils.CLIENT_NAME,
164+
() -> {
165+
writer.write(
166+
"page = await makePagedClientRequest(config.client, input, ...additionalArguments);");
167+
}
168+
);
169+
writer.openBlock("else {", "}", () -> {
170+
writer.write("throw new Error(\"Invalid client, expected $L | $L\");",
171+
DocumentClientUtils.CLIENT_FULL_NAME, DocumentClientUtils.CLIENT_NAME);
172+
});
173+
174+
writer.write("yield page;");
175+
writer.write("token = page$L;", destructurePath(outputTokenName));
176+
177+
writer.write("hasNext = !!(token);");
178+
});
179+
180+
writer.write("// @ts-ignore");
181+
writer.write("return undefined;");
182+
});
183+
}
184+
185+
186+
/**
187+
* Paginated command that calls client.method({...}) under the hood. This is meant for server side environments and
188+
* exposes the entire service.
189+
*/
190+
private void writeMethodRequest() {
191+
writer.writeDocs("@private");
192+
writer.openBlock(
193+
"const makePagedRequest = async (client: $L, input: $L, ...args: any): Promise<$L> => {",
194+
"}", DocumentClientUtils.CLIENT_FULL_NAME, inputTypeName,
195+
outputTypeName, () -> {
196+
writer.write("// @ts-ignore");
197+
writer.write("return await client.$L(input, ...args);", methodName);
198+
});
199+
}
200+
201+
/**
202+
* Paginated command that calls CommandClient().send({...}) under the hood. This is meant for client side (browser)
203+
* environments and does not generally expose the entire service.
204+
*/
205+
private void writeCommandRequest() {
206+
writer.writeDocs("@private");
207+
writer.openBlock(
208+
"const makePagedClientRequest = async (client: $L, input: $L, ...args: any): Promise<$L> => {",
209+
"}", DocumentClientUtils.CLIENT_NAME, inputTypeName,
210+
outputTypeName, () -> {
211+
writer.write("// @ts-ignore");
212+
writer.write("return await client.send(new $L(input), ...args);", operationTypeName);
213+
});
214+
}
215+
}

‎codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/DocumentClientUtils.java

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ final class DocumentClientUtils {
3838
static final String CLIENT_CONFIG_NAME = getResolvedConfigTypeName(CLIENT_NAME);
3939
static final String CLIENT_COMMANDS_FOLDER = "commands";
4040
static final String CLIENT_UTILS_FILE = "utils";
41+
static final String DOC_CLIENT_PREFIX = "doc-client-";
4142

4243
static final String CLIENT_TRANSLATE_CONFIG_KEY = "translateConfig";
4344
static final String CLIENT_TRANSLATE_CONFIG_TYPE = "TranslateConfig";

‎lib/lib-dynamodb/src/index.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
11
export * from "./DynamoDBDocument";
22
export * from "./DynamoDBDocumentClient";
33
export * from "./commands";
4+
export * from "./pagination";

‎lib/lib-dynamodb/src/pagination/Interfaces.ts

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
import { PaginationConfiguration } from "@aws-sdk/types";
2+
3+
import { DynamoDBDocument } from "../DynamoDBDocument";
4+
import { DynamoDBDocumentClient } from "../DynamoDBDocumentClient";
5+
6+
export interface DynamoDBDocumentPaginationConfiguration extends PaginationConfiguration {
7+
client: DynamoDBDocument | DynamoDBDocumentClient;
8+
}

‎lib/lib-dynamodb/src/pagination/QueryPaginator.ts

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
import { Paginator } from "@aws-sdk/types";
2+
3+
import { QueryCommand, QueryCommandInput, QueryCommandOutput } from "../commands/QueryCommand";
4+
import { DynamoDBDocument } from "../DynamoDBDocument";
5+
import { DynamoDBDocumentClient } from "../DynamoDBDocumentClient";
6+
import { DynamoDBDocumentPaginationConfiguration } from "./Interfaces";
7+
8+
/**
9+
* @private
10+
*/
11+
const makePagedClientRequest = async (
12+
client: DynamoDBDocumentClient,
13+
input: QueryCommandInput,
14+
...args: any
15+
): Promise<QueryCommandOutput> => {
16+
// @ts-ignore
17+
return await client.send(new QueryCommand(input), ...args);
18+
};
19+
/**
20+
* @private
21+
*/
22+
const makePagedRequest = async (
23+
client: DynamoDBDocument,
24+
input: QueryCommandInput,
25+
...args: any
26+
): Promise<QueryCommandOutput> => {
27+
// @ts-ignore
28+
return await client.query(input, ...args);
29+
};
30+
export async function* paginateQuery(
31+
config: DynamoDBDocumentPaginationConfiguration,
32+
input: QueryCommandInput,
33+
...additionalArguments: any
34+
): Paginator<QueryCommandOutput> {
35+
// ToDo: replace with actual type instead of typeof input.ExclusiveStartKey
36+
let token: typeof input.ExclusiveStartKey | undefined = config.startingToken || undefined;
37+
let hasNext = true;
38+
let page: QueryCommandOutput;
39+
while (hasNext) {
40+
input.ExclusiveStartKey = token;
41+
input["Limit"] = config.pageSize;
42+
if (config.client instanceof DynamoDBDocument) {
43+
page = await makePagedRequest(config.client, input, ...additionalArguments);
44+
} else if (config.client instanceof DynamoDBDocumentClient) {
45+
page = await makePagedClientRequest(config.client, input, ...additionalArguments);
46+
} else {
47+
throw new Error("Invalid client, expected DynamoDBDocument | DynamoDBDocumentClient");
48+
}
49+
yield page;
50+
token = page.LastEvaluatedKey;
51+
hasNext = !!token;
52+
}
53+
// @ts-ignore
54+
return undefined;
55+
}

0 commit comments

Comments
 (0)