feat(pagination): add middleware support (#1480)

* feat(pagination): add middleware support * use command callback * formatting * update changeset --------- Co-authored-by: George Fu <[email protected]>
smithy-lang · Jan 9, 2025 · 2aff9df · 2aff9df
1 parent 292c134
commit 2aff9df
Show file tree

Hide file tree

Showing 4 changed files with 91 additions and 5 deletions.
diff --git a/.changeset/fast-cups-speak.md 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
@@ -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<PaginationConfiguration, { inToken?: string }, { outToken: string }>(
       Client,
@@ -112,4 +124,56 @@ describe(createPaginator.name, () => {
 
     expect(pages).toEqual(5);
   });
+
+  it("should allow modification of the instantiated command", async () => {
+    const paginate = createPaginator<PaginationConfiguration, { inToken?: string }, { outToken: string }>(
+      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<typeof paginate>[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
@@ -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 <ClientType extends Client<any, any, any>,
   CommandCtor: any,
   client: ClientType,
   input: InputType,
+  withCommand: (command: Command<any, any, any, any, any>) => typeof command = (_) => _,
   ...args: any[]
 ): Promise<OutputType> => {
-  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
@@ -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<any, any, any, any, any>) => typeof command;
 }