Skip to content

Commit

Permalink
feat(lib-dynamodb): add pagination (#3069)
Browse files Browse the repository at this point in the history
  • Loading branch information
trivikr committed Nov 29, 2021
1 parent 31f3478 commit 51480df
Show file tree
Hide file tree
Showing 8 changed files with 375 additions and 9 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
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;
import software.amazon.smithy.typescript.codegen.TypeScriptSettings;
import software.amazon.smithy.typescript.codegen.TypeScriptWriter;
import software.amazon.smithy.typescript.codegen.integration.TypeScriptIntegration;
Expand All @@ -49,14 +50,13 @@ public void writeAdditionalFiles(
) {
ServiceShape service = settings.getService(model);
if (testServiceId(service, "DynamoDB")) {
String docClientPrefix = "doc-client-";
Set<OperationShape> containedOperations =
new TreeSet<>(TopDownIndex.of(model).getContainedOperations(service));
List<OperationShape> overridenOperationsList = new ArrayList<>();

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

if (DocumentClientUtils.containsAttributeValue(model, symbolProvider, operation)) {
Expand All @@ -65,16 +65,25 @@ public void writeAdditionalFiles(
writer -> new DocumentClientCommandGenerator(
settings, model, operation, symbolProvider, writer).run()
);

if (operation.hasTrait(PaginatedTrait.ID)) {
String paginationFileName = DocumentClientPaginationGenerator.getOutputFilelocation(operation);
writerFactory.accept(paginationFileName, paginationWriter ->
new DocumentClientPaginationGenerator(model, service, operation, symbolProvider,
paginationWriter).run());
}
}
}

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

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

writerFactory.accept(String.format("%s%s/index.ts", docClientPrefix,
writerFactory.accept(String.format("%s%s/index.ts", DocumentClientUtils.DOC_CLIENT_PREFIX,
DocumentClientUtils.CLIENT_COMMANDS_FOLDER), writer -> {
for (OperationShape operation : overridenOperationsList) {
String operationFileName = DocumentClientUtils.getModifiedName(
Expand All @@ -83,19 +92,38 @@ public void writeAdditionalFiles(
writer.write("export * from './$L';", operationFileName);
}
});
writerFactory.accept(String.format("%sindex.ts", docClientPrefix), writer -> {

writerFactory.accept(String.format("%s%s/index.ts", DocumentClientUtils.DOC_CLIENT_PREFIX,
DocumentClientPaginationGenerator.PAGINATION_FOLDER), writer -> {
writer.write("export * from './Interfaces';");
for (OperationShape operation : overridenOperationsList) {
if (operation.hasTrait(PaginatedTrait.ID)) {
String paginationFileName =
DocumentClientUtils.getModifiedName(operation.getId().getName()) + "Paginator";
writer.write("export * from './$L';", paginationFileName);
}
}
});

String paginationInterfaceFileName = DocumentClientPaginationGenerator.getInterfaceFilelocation();
writerFactory.accept(paginationInterfaceFileName, paginationWriter ->
DocumentClientPaginationGenerator.generateServicePaginationInterfaces(paginationWriter));

writerFactory.accept(String.format("%sindex.ts", DocumentClientUtils.DOC_CLIENT_PREFIX), writer -> {
writer.write("export * from './commands';");
writer.write("export * from './pagination';");
writer.write("export * from './$L';", DocumentClientUtils.CLIENT_NAME);
writer.write("export * from './$L';", DocumentClientUtils.CLIENT_FULL_NAME);
});

String utilsFileLocation = String.format("%s%s", docClientPrefix, DocumentClientUtils.CLIENT_UTILS_FILE);
writerFactory.accept(String.format("%s%s/%s.ts", docClientPrefix,
String utilsFileLocation = String.format("%s%s", DocumentClientUtils.DOC_CLIENT_PREFIX,
DocumentClientUtils.CLIENT_UTILS_FILE);
writerFactory.accept(String.format("%s%s/%s.ts", DocumentClientUtils.DOC_CLIENT_PREFIX,
DocumentClientUtils.CLIENT_COMMANDS_FOLDER, DocumentClientUtils.CLIENT_UTILS_FILE), writer -> {
writer.write(IoUtils.readUtf8Resource(AddDocumentClientPlugin.class,
String.format("%s.ts", utilsFileLocation)));
});
writerFactory.accept(String.format("%s%s/%s.spec.ts", docClientPrefix,
writerFactory.accept(String.format("%s%s/%s.spec.ts", DocumentClientUtils.DOC_CLIENT_PREFIX,
DocumentClientUtils.CLIENT_COMMANDS_FOLDER, DocumentClientUtils.CLIENT_UTILS_FILE), writer -> {
writer.write(IoUtils.readUtf8Resource(AddDocumentClientPlugin.class,
String.format("%s.spec.ts", utilsFileLocation)));
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,215 @@
/*
* Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License").
* You may not use this file except in compliance with the License.
* A copy of the License is located at
*
* http://aws.amazon.com/apache2.0
*
* or in the "license" file accompanying this file. This file is distributed
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
* express or implied. See the License for the specific language governing
* permissions and limitations under the License.
*/

package software.amazon.smithy.aws.typescript.codegen;

import java.nio.file.Paths;
import java.util.Optional;
import software.amazon.smithy.codegen.core.CodegenException;
import software.amazon.smithy.codegen.core.Symbol;
import software.amazon.smithy.codegen.core.SymbolProvider;
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.OperationShape;
import software.amazon.smithy.model.shapes.ServiceShape;
import software.amazon.smithy.typescript.codegen.TypeScriptWriter;
import software.amazon.smithy.utils.SmithyInternalApi;

@SmithyInternalApi
final class DocumentClientPaginationGenerator implements Runnable {

static final String PAGINATION_FOLDER = "pagination";

private final TypeScriptWriter writer;
private final PaginationInfo paginatedInfo;

private final String operationTypeName;
private final String inputTypeName;
private final String outputTypeName;

private final String operationName;
private final String methodName;
private final String paginationType;

DocumentClientPaginationGenerator(
Model model,
ServiceShape service,
OperationShape operation,
SymbolProvider symbolProvider,
TypeScriptWriter writer
) {
this.writer = writer;

Symbol operationSymbol = symbolProvider.toSymbol(operation);
Symbol inputSymbol = symbolProvider.toSymbol(operation).expectProperty("inputType", Symbol.class);
Symbol outputSymbol = symbolProvider.toSymbol(operation).expectProperty("outputType", Symbol.class);

this.operationTypeName = DocumentClientUtils.getModifiedName(operationSymbol.getName());
this.inputTypeName = DocumentClientUtils.getModifiedName(inputSymbol.getName());
this.outputTypeName = DocumentClientUtils.getModifiedName(outputSymbol.getName());

// e.g. listObjects
this.operationName = operationTypeName.replace("Command", "");
this.methodName = Character.toLowerCase(operationName.charAt(0)) + operationName.substring(1);
this.paginationType = DocumentClientUtils.CLIENT_FULL_NAME + "PaginationConfiguration";

PaginatedIndex paginatedIndex = PaginatedIndex.of(model);
Optional<PaginationInfo> paginationInfo = paginatedIndex.getPaginationInfo(service, operation);
this.paginatedInfo = paginationInfo.orElseThrow(() -> {
return new CodegenException("Expected Paginator to have pagination information.");
});
}

@Override
public void run() {
// Import Service Types
String commandFileLocation = Paths.get(".", DocumentClientUtils.CLIENT_COMMANDS_FOLDER,
DocumentClientUtils.getModifiedName(operationTypeName)).toString();
writer.addImport(operationTypeName, operationTypeName, commandFileLocation);
writer.addImport(inputTypeName, inputTypeName, commandFileLocation);
writer.addImport(outputTypeName, outputTypeName, commandFileLocation);
writer.addImport(
DocumentClientUtils.CLIENT_NAME,
DocumentClientUtils.CLIENT_NAME,
Paths.get(".", DocumentClientUtils.CLIENT_NAME).toString());
writer.addImport(
DocumentClientUtils.CLIENT_FULL_NAME,
DocumentClientUtils.CLIENT_FULL_NAME,
Paths.get(".", DocumentClientUtils.CLIENT_FULL_NAME).toString());

// Import Pagination types
writer.addImport("Paginator", "Paginator", "@aws-sdk/types");
writer.addImport(paginationType, paginationType,
Paths.get(".", getInterfaceFilelocation().replace(".ts", "")).toString());

writeCommandRequest();
writeMethodRequest();
writePager();
}

static String getOutputFilelocation(OperationShape operation) {
return String.format("%s%s/%s.ts", DocumentClientUtils.DOC_CLIENT_PREFIX,
DocumentClientPaginationGenerator.PAGINATION_FOLDER,
DocumentClientUtils.getModifiedName(operation.getId().getName()) + "Paginator");
}

static String getInterfaceFilelocation() {
return String.format("%s%s/%s.ts", DocumentClientUtils.DOC_CLIENT_PREFIX,
DocumentClientPaginationGenerator.PAGINATION_FOLDER, "Interfaces");
}

static void generateServicePaginationInterfaces(TypeScriptWriter writer) {
writer.addImport("PaginationConfiguration", "PaginationConfiguration", "@aws-sdk/types");

writer.addImport(
DocumentClientUtils.CLIENT_NAME,
DocumentClientUtils.CLIENT_NAME,
Paths.get(".", DocumentClientUtils.CLIENT_NAME).toString());
writer.addImport(
DocumentClientUtils.CLIENT_FULL_NAME,
DocumentClientUtils.CLIENT_FULL_NAME,
Paths.get(".", DocumentClientUtils.CLIENT_FULL_NAME).toString());

writer.openBlock("export interface $LPaginationConfiguration extends PaginationConfiguration {",
"}", DocumentClientUtils.CLIENT_FULL_NAME, () -> {
writer.write("client: $L | $L;", DocumentClientUtils.CLIENT_FULL_NAME, DocumentClientUtils.CLIENT_NAME);
});
}

private String destructurePath(String path) {
return "." + path.replace(".", "!.");
}

private void writePager() {
String inputTokenName = paginatedInfo.getPaginatedTrait().getInputToken().get();
String outputTokenName = paginatedInfo.getPaginatedTrait().getOutputToken().get();

writer.openBlock(
"export async function* paginate$L(config: $L, input: $L, ...additionalArguments: any): Paginator<$L>{",
"}", operationName, paginationType, inputTypeName, outputTypeName, () -> {
String destructuredInputTokenName = destructurePath(inputTokenName);
writer.write("// ToDo: replace with actual type instead of typeof input$L", destructuredInputTokenName);
writer.write("let token: typeof input$L | undefined = config.startingToken || undefined;",
destructuredInputTokenName);

writer.write("let hasNext = true;");
writer.write("let page: $L;", outputTypeName);
writer.openBlock("while (hasNext) {", "}", () -> {
writer.write("input$L = token;", destructuredInputTokenName);

if (paginatedInfo.getPageSizeMember().isPresent()) {
String pageSize = paginatedInfo.getPageSizeMember().get().getMemberName();
writer.write("input[$S] = config.pageSize;", pageSize);
}

writer.openBlock("if (config.client instanceof $L) {", "}", DocumentClientUtils.CLIENT_FULL_NAME,
() -> {
writer.write("page = await makePagedRequest(config.client, input, ...additionalArguments);");
}
);
writer.openBlock("else if (config.client instanceof $L) {", "}", DocumentClientUtils.CLIENT_NAME,
() -> {
writer.write(
"page = await makePagedClientRequest(config.client, input, ...additionalArguments);");
}
);
writer.openBlock("else {", "}", () -> {
writer.write("throw new Error(\"Invalid client, expected $L | $L\");",
DocumentClientUtils.CLIENT_FULL_NAME, DocumentClientUtils.CLIENT_NAME);
});

writer.write("yield page;");
writer.write("token = page$L;", destructurePath(outputTokenName));

writer.write("hasNext = !!(token);");
});

writer.write("// @ts-ignore");
writer.write("return undefined;");
});
}


/**
* Paginated command that calls client.method({...}) under the hood. This is meant for server side environments and
* exposes the entire service.
*/
private void writeMethodRequest() {
writer.writeDocs("@private");
writer.openBlock(
"const makePagedRequest = async (client: $L, input: $L, ...args: any): Promise<$L> => {",
"}", DocumentClientUtils.CLIENT_FULL_NAME, inputTypeName,
outputTypeName, () -> {
writer.write("// @ts-ignore");
writer.write("return await client.$L(input, ...args);", methodName);
});
}

/**
* Paginated command that calls CommandClient().send({...}) under the hood. This is meant for client side (browser)
* environments and does not generally expose the entire service.
*/
private void writeCommandRequest() {
writer.writeDocs("@private");
writer.openBlock(
"const makePagedClientRequest = async (client: $L, input: $L, ...args: any): Promise<$L> => {",
"}", DocumentClientUtils.CLIENT_NAME, inputTypeName,
outputTypeName, () -> {
writer.write("// @ts-ignore");
writer.write("return await client.send(new $L(input), ...args);", operationTypeName);
});
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ final class DocumentClientUtils {
static final String CLIENT_CONFIG_NAME = getResolvedConfigTypeName(CLIENT_NAME);
static final String CLIENT_COMMANDS_FOLDER = "commands";
static final String CLIENT_UTILS_FILE = "utils";
static final String DOC_CLIENT_PREFIX = "doc-client-";

static final String CLIENT_TRANSLATE_CONFIG_KEY = "translateConfig";
static final String CLIENT_TRANSLATE_CONFIG_TYPE = "TranslateConfig";
Expand Down
1 change: 1 addition & 0 deletions lib/lib-dynamodb/src/index.ts
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
export * from "./DynamoDBDocument";
export * from "./DynamoDBDocumentClient";
export * from "./commands";
export * from "./pagination";
8 changes: 8 additions & 0 deletions lib/lib-dynamodb/src/pagination/Interfaces.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
import { PaginationConfiguration } from "@aws-sdk/types";

import { DynamoDBDocument } from "../DynamoDBDocument";
import { DynamoDBDocumentClient } from "../DynamoDBDocumentClient";

export interface DynamoDBDocumentPaginationConfiguration extends PaginationConfiguration {
client: DynamoDBDocument | DynamoDBDocumentClient;
}
55 changes: 55 additions & 0 deletions lib/lib-dynamodb/src/pagination/QueryPaginator.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
import { Paginator } from "@aws-sdk/types";

import { QueryCommand, QueryCommandInput, QueryCommandOutput } from "../commands/QueryCommand";
import { DynamoDBDocument } from "../DynamoDBDocument";
import { DynamoDBDocumentClient } from "../DynamoDBDocumentClient";
import { DynamoDBDocumentPaginationConfiguration } from "./Interfaces";

/**
* @private
*/
const makePagedClientRequest = async (
client: DynamoDBDocumentClient,
input: QueryCommandInput,
...args: any
): Promise<QueryCommandOutput> => {
// @ts-ignore
return await client.send(new QueryCommand(input), ...args);
};
/**
* @private
*/
const makePagedRequest = async (
client: DynamoDBDocument,
input: QueryCommandInput,
...args: any
): Promise<QueryCommandOutput> => {
// @ts-ignore
return await client.query(input, ...args);
};
export async function* paginateQuery(
config: DynamoDBDocumentPaginationConfiguration,
input: QueryCommandInput,
...additionalArguments: any
): Paginator<QueryCommandOutput> {
// ToDo: replace with actual type instead of typeof input.ExclusiveStartKey
let token: typeof input.ExclusiveStartKey | undefined = config.startingToken || undefined;
let hasNext = true;
let page: QueryCommandOutput;
while (hasNext) {
input.ExclusiveStartKey = token;
input["Limit"] = config.pageSize;
if (config.client instanceof DynamoDBDocument) {
page = await makePagedRequest(config.client, input, ...additionalArguments);
} else if (config.client instanceof DynamoDBDocumentClient) {
page = await makePagedClientRequest(config.client, input, ...additionalArguments);
} else {
throw new Error("Invalid client, expected DynamoDBDocument | DynamoDBDocumentClient");
}
yield page;
token = page.LastEvaluatedKey;
hasNext = !!token;
}
// @ts-ignore
return undefined;
}

0 comments on commit 51480df

Please sign in to comment.