diff --git a/.changeset/fast-cups-speak.md b/.changeset/fast-cups-speak.md new file mode 100644 index 00000000000..82d8c5a14c0 --- /dev/null +++ b/.changeset/fast-cups-speak.md @@ -0,0 +1,6 @@ +--- +"@smithy/types": minor +"@smithy/core": minor +--- + +Added middleware support to pagination diff --git a/packages/core/src/pagination/createPaginator.spec.ts b/packages/core/src/pagination/createPaginator.spec.ts index 6ab8085fc35..808ae5093e0 100644 --- a/packages/core/src/pagination/createPaginator.spec.ts +++ b/packages/core/src/pagination/createPaginator.spec.ts @@ -1,5 +1,5 @@ import { PaginationConfiguration } from "@smithy/types"; -import { describe, expect, test as it } from "vitest"; +import { afterEach, describe, expect, test as it, vi } from "vitest"; import { createPaginator } from "./createPaginator"; @@ -20,6 +20,10 @@ describe(createPaginator.name, () => { } } class CommandObjectToken { + public middlewareStack = { + add: vi.fn(), + addRelativeTo: vi.fn(), + }; public constructor(public input: any) { expect(input).toEqual({ sizeToken: 100, @@ -33,6 +37,10 @@ describe(createPaginator.name, () => { } class CommandStringToken { + public middlewareStack = { + add: vi.fn(), + addRelativeTo: vi.fn(), + }; public constructor(public input: any) { expect(input).toEqual({ sizeToken: 100, @@ -41,6 +49,10 @@ describe(createPaginator.name, () => { } } + afterEach(() => { + vi.resetAllMocks(); + }); + it("should create a paginator", async () => { const paginate = createPaginator( Client, @@ -112,4 +124,56 @@ describe(createPaginator.name, () => { expect(pages).toEqual(5); }); + + it("should allow modification of the instantiated command", async () => { + const paginate = createPaginator( + Client, + CommandObjectToken, + "inToken", + "outToken", + "sizeToken" + ); + + let pages = 0; + const client: any = new Client(); + vi.spyOn(client, "send"); + const config = { + client, + pageSize: 100, + startingToken: { + outToken2: { + outToken3: "TOKEN_VALUE", + }, + }, + withCommand(command) { + command.middlewareStack.add((next) => (args) => next(args)); + command.middlewareStack.addRelativeTo((next: any) => (args: any) => next(args), { + toMiddleware: "", + relation: "before", + }); + expect(command.middlewareStack.add).toHaveBeenCalledTimes(1); + expect(command.middlewareStack.addRelativeTo).toHaveBeenCalledTimes(1); + return command; + }, + } as Parameters[0]; + vi.spyOn(config, "withCommand"); + + for await (const page of paginate(config, {})) { + pages += 1; + if (pages === 5) { + expect(page.outToken).toBeUndefined(); + } else { + expect(page.outToken).toEqual({ + outToken2: { + outToken3: "TOKEN_VALUE", + }, + }); + } + } + + expect(pages).toEqual(5); + expect(client.send).toHaveBeenCalledTimes(5); + expect(config.withCommand).toHaveBeenCalledTimes(5); + expect(config.withCommand).toHaveBeenCalledWith(expect.any(CommandObjectToken)); + }); }); diff --git a/packages/core/src/pagination/createPaginator.ts b/packages/core/src/pagination/createPaginator.ts index dd3fb09adcb..1f762be9782 100644 --- a/packages/core/src/pagination/createPaginator.ts +++ b/packages/core/src/pagination/createPaginator.ts @@ -1,4 +1,4 @@ -import type { Client, PaginationConfiguration, Paginator } from "@smithy/types"; +import type { Client, Command, PaginationConfiguration, Paginator } from "@smithy/types"; /** * @internal @@ -7,9 +7,12 @@ const makePagedClientRequest = async , CommandCtor: any, client: ClientType, input: InputType, + withCommand: (command: Command) => typeof command = (_) => _, ...args: any[] ): Promise => { - return await client.send(new CommandCtor(input), ...args); + let command = new CommandCtor(input); + command = withCommand(command); + return await client.send(command, ...args); }; /** @@ -43,7 +46,13 @@ export function createPaginator< (input as any)[pageSizeTokenName] = (input as any)[pageSizeTokenName] ?? config.pageSize; } if (config.client instanceof ClientCtor) { - page = await makePagedClientRequest(CommandCtor, config.client, input, ...additionalArguments); + page = await makePagedClientRequest( + CommandCtor, + config.client, + input, + config.withCommand, + ...additionalArguments + ); } else { throw new Error(`Invalid client, expected instance of ${ClientCtor.name}`); } diff --git a/packages/types/src/pagination.ts b/packages/types/src/pagination.ts index 07eb21f5a86..624441dcab0 100644 --- a/packages/types/src/pagination.ts +++ b/packages/types/src/pagination.ts @@ -1,4 +1,5 @@ -import { Client } from "./client"; +import type { Client } from "./client"; +import type { Command } from "./command"; /** * @public @@ -25,4 +26,10 @@ export interface PaginationConfiguration { * instead of when it is not present. */ stopOnSameToken?: boolean; + /** + * @param command - reference to the instantiated command. This callback is executed + * prior to sending the command with the paginator's client. + * @returns the original command or a replacement. + */ + withCommand?: (command: Command) => typeof command; }