fix(pagination): allow operation input to specify pagination token (#1500)

* fix(pagination): only use startingToken if it exists

* add unit tests

* add changeset

* undo types change

* default return value of withCommand in paginator

---------

Co-authored-by: George Fu <[email protected]>

Author: icholy

.changeset/lazy-geese-brush.md

@@ -0,0 +1,6 @@
+---
+"@smithy/types": minor
+"@smithy/core": minor
+---
+
+allow paginator token fallback to be specified by operation input

packages/core/src/pagination/createPaginator.spec.ts

@@ -36,6 +36,17 @@ describe(createPaginator.name, () => {
     }
   }
 
+  class ClientStringToken {
+    private pages = 5;
+    async send(command: any) {
+      if (--this.pages > 0) {
+        return {
+          outToken: command.input.inToken,
+        };
+      }
+      return {};
+    }
+  }
   class CommandStringToken {
     public middlewareStack = {
       add: vi.fn(),
@@ -91,6 +102,158 @@ describe(createPaginator.name, () => {
     expect(pages).toEqual(5);
   });
 
+  it("should prioritize token set in paginator config, fallback to token set in input parameters", async () => {
+    class CommandExpectPaginatorConfigToken {
+      public constructor(public input: any) {
+        expect(input).toMatchObject({
+          inToken: "abc",
+        });
+      }
+    }
+    class CommandExpectOperationInputToken {
+      public constructor(public input: any) {
+        expect(input).toMatchObject({
+          inToken: "xyz",
+        });
+      }
+    }
+    {
+      const paginate = createPaginator<
+        PaginationConfiguration,
+        { inToken?: string; sizeToken?: number },
+        { outToken: string }
+      >(ClientStringToken, CommandExpectPaginatorConfigToken, "inToken", "outToken", "sizeToken");
+
+      let pages = 0;
+      const client = new ClientStringToken() as any;
+
+      for await (const page of paginate(
+        {
+          client,
+          startingToken: "abc",
+        },
+        {
+          inToken: "xyz",
+        }
+      )) {
+        pages += 1;
+        expect(page).toBeDefined();
+      }
+
+      expect(pages).toEqual(5);
+    }
+    {
+      const paginate = createPaginator<
+        PaginationConfiguration,
+        { inToken?: string; sizeToken?: number },
+        { outToken: string }
+      >(ClientStringToken, CommandExpectOperationInputToken, "inToken", "outToken", "sizeToken");
+
+      let pages = 0;
+      const client = new ClientStringToken() as any;
+
+      for await (const page of paginate(
+        {
+          client,
+        },
+        {
+          inToken: "xyz",
+        }
+      )) {
+        pages += 1;
+        expect(page).toBeDefined();
+      }
+
+      expect(pages).toEqual(5);
+    }
+  });
+
+  it("should prioritize page size set in operation input, fallback to page size set in paginator config (inverted from token priority)", async () => {
+    class CommandExpectPaginatorPageSize {
+      public constructor(public input: any) {
+        expect(input).toMatchObject({
+          sizeToken: 100,
+        });
+      }
+    }
+    class CommandExpectOperationInputPageSize {
+      public constructor(public input: any) {
+        expect(input).toMatchObject({
+          sizeToken: 99,
+        });
+      }
+    }
+    {
+      const paginate = createPaginator<
+        PaginationConfiguration,
+        { inToken?: string; sizeToken?: number },
+        { outToken: string }
+      >(ClientStringToken, CommandExpectPaginatorPageSize, "inToken", "outToken", "sizeToken");
+
+      let pages = 0;
+      const client = new ClientStringToken() as any;
+
+      for await (const page of paginate(
+        {
+          client,
+          pageSize: 100,
+        },
+        {
+          inToken: "abc",
+        }
+      )) {
+        pages += 1;
+        expect(page).toBeDefined();
+      }
+
+      expect(pages).toEqual(5);
+    }
+    {
+      const paginate = createPaginator<
+        PaginationConfiguration,
+        { inToken?: string; sizeToken?: number },
+        { outToken: string }
+      >(ClientStringToken, CommandExpectOperationInputPageSize, "inToken", "outToken", "sizeToken");
+
+      let pages = 0;
+      const client = new ClientStringToken() as any;
+
+      for await (const page of paginate(
+        {
+          client,
+          pageSize: 100,
+        },
+        {
+          sizeToken: 99,
+          inToken: "abc",
+        }
+      )) {
+        pages += 1;
+        expect(page).toBeDefined();
+      }
+
+      expect(pages).toEqual(5);
+    }
+  });
+
+  it("should have the correct AsyncGenerator.TNext type", async () => {
+    const paginate = createPaginator<
+      PaginationConfiguration,
+      { inToken?: string; sizeToken: number },
+      {
+        outToken: string;
+      }
+    >(ClientStringToken, CommandStringToken, "inToken", "outToken.outToken2.outToken3", "sizeToken");
+    const asyncGenerator = paginate(
+      { client: new ClientStringToken() as any },
+      { inToken: "TOKEN_VALUE", sizeToken: 100 }
+    );
+
+    const { value, done } = await asyncGenerator.next();
+    expect(value?.outToken).toBeTypeOf("string");
+    expect(done).toBe(false);
+  });
+
   it("should handle deep paths", async () => {
     const paginate = createPaginator<
       PaginationConfiguration,

packages/core/src/pagination/createPaginator.ts

@@ -11,7 +11,7 @@ const makePagedClientRequest = async <ClientType extends Client<any, any, any>,
   ...args: any[]
 ): Promise<OutputType> => {
   let command = new CommandCtor(input);
-  command = withCommand(command);
+  command = withCommand(command) ?? command;
   return await client.send(command, ...args);
 };
 
@@ -36,14 +36,16 @@ export function createPaginator<
     input: InputType,
     ...additionalArguments: any[]
   ): Paginator<OutputType> {
-    let token: any = config.startingToken || undefined;
+    const _input = input as any;
+    // for legacy reasons this coalescing order is inverted from that of pageSize.
+    let token: any = config.startingToken ?? _input[inputTokenName];
     let hasNext = true;
     let page: OutputType;
 
     while (hasNext) {
-      (input as any)[inputTokenName] = token;
+      _input[inputTokenName] = token;
       if (pageSizeTokenName) {
-        (input as any)[pageSizeTokenName] = (input as any)[pageSizeTokenName] ?? config.pageSize;
+        _input[pageSizeTokenName] = _input[pageSizeTokenName] ?? config.pageSize;
       }
       if (config.client instanceof ClientCtor) {
         page = await makePagedClientRequest(
@@ -61,7 +63,6 @@ export function createPaginator<
       token = get(page, outputTokenName);
       hasNext = !!(token && (!config.stopOnSameToken || token !== prevToken));
     }
-    // @ts-ignore
     return undefined;
   };
 }

packages/types/src/pagination.ts

@@ -29,7 +29,7 @@ export interface PaginationConfiguration {
   /**
    * @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.
+   * @returns the original command or a replacement, defaulting to the original command object.
    */
-  withCommand?: (command: Command<any, any, any, any, any>) => typeof command;
+  withCommand?: (command: Command<any, any, any, any, any>) => typeof command | undefined;
 }