From 265f2ad9fd586e6c33793ff6cb47bb9174c52cdb Mon Sep 17 00:00:00 2001
From: George Fu <kuhe@users.noreply.github.com>
Date: Wed, 7 Feb 2024 15:04:45 +0000
Subject: [PATCH] fix: handle multi-part token paths in paginator

---
 .changeset/four-steaks-sip.md                 |   5 +
 .../src/pagination/createPaginator.spec.ts    | 114 ++++++++++++++++++
 .../core/src/pagination/createPaginator.ts    |  17 ++-
 3 files changed, 135 insertions(+), 1 deletion(-)
 create mode 100644 .changeset/four-steaks-sip.md
 create mode 100644 packages/core/src/pagination/createPaginator.spec.ts

diff --git a/.changeset/four-steaks-sip.md b/.changeset/four-steaks-sip.md
new file mode 100644
index 00000000000..faadac685b6
--- /dev/null
+++ b/.changeset/four-steaks-sip.md
@@ -0,0 +1,5 @@
+---
+"@smithy/core": patch
+---
+
+handle multi-part input token in paginator
diff --git a/packages/core/src/pagination/createPaginator.spec.ts b/packages/core/src/pagination/createPaginator.spec.ts
new file mode 100644
index 00000000000..ade5107b899
--- /dev/null
+++ b/packages/core/src/pagination/createPaginator.spec.ts
@@ -0,0 +1,114 @@
+import { PaginationConfiguration } from "@smithy/types";
+
+import { createPaginator } from "./createPaginator";
+
+describe(createPaginator.name, () => {
+  class Client {
+    private pages = 5;
+    async send(command: CommandObjectToken | CommandStringToken) {
+      if (--this.pages > 0) {
+        return {
+          outToken: {
+            outToken2: {
+              outToken3: "TOKEN_VALUE",
+            },
+          },
+        };
+      }
+      return {};
+    }
+  }
+  class CommandObjectToken {
+    public constructor(public input: any) {
+      expect(input).toEqual({
+        sizeToken: 100,
+        inToken: {
+          outToken2: {
+            outToken3: "TOKEN_VALUE",
+          },
+        },
+      });
+    }
+  }
+
+  class CommandStringToken {
+    public constructor(public input: any) {
+      expect(input).toEqual({
+        sizeToken: 100,
+        inToken: "TOKEN_VALUE",
+      });
+    }
+  }
+
+  it("should create a paginator", async () => {
+    const paginate = createPaginator<PaginationConfiguration, { inToken?: string }, { outToken: string }>(
+      Client,
+      CommandObjectToken,
+      "inToken",
+      "outToken",
+      "sizeToken"
+    );
+
+    let pages = 0;
+
+    for await (const page of paginate(
+      {
+        client: new Client() as any,
+        pageSize: 100,
+        startingToken: {
+          outToken2: {
+            outToken3: "TOKEN_VALUE",
+          },
+        },
+      },
+      {}
+    )) {
+      pages += 1;
+      if (pages === 5) {
+        expect(page.outToken).toBeUndefined();
+      } else {
+        expect(page.outToken).toEqual({
+          outToken2: {
+            outToken3: "TOKEN_VALUE",
+          },
+        });
+      }
+    }
+
+    expect(pages).toEqual(5);
+  });
+
+  it("should handle deep paths", async () => {
+    const paginate = createPaginator<
+      PaginationConfiguration,
+      { inToken?: string },
+      {
+        outToken: {
+          outToken2: {
+            outToken3: string;
+          };
+        };
+      }
+    >(Client, CommandStringToken, "inToken", "outToken.outToken2.outToken3", "sizeToken");
+
+    let pages = 0;
+
+    for await (const page of paginate(
+      {
+        client: new Client() as any,
+        pageSize: 100,
+        startingToken: "TOKEN_VALUE",
+      },
+      {}
+    )) {
+      pages += 1;
+      if (pages === 5) {
+        expect(page.outToken).toBeUndefined();
+      } else {
+        expect(page.outToken.outToken2.outToken3).toEqual("TOKEN_VALUE");
+      }
+    }
+
+    expect(pages).toEqual(5);
+  });
+});
diff --git a/packages/core/src/pagination/createPaginator.ts b/packages/core/src/pagination/createPaginator.ts
index 5421528d5b7..a82a49209b8 100644
--- a/packages/core/src/pagination/createPaginator.ts
+++ b/packages/core/src/pagination/createPaginator.ts
@@ -49,10 +49,25 @@ export function createPaginator<
       }
       yield page;
       const prevToken = token;
-      token = (page as any)[outputTokenName];
+      token = get(page, outputTokenName);
       hasNext = !!(token && (!config.stopOnSameToken || token !== prevToken));
     }
     // @ts-ignore
     return undefined;
   };
 }
+
+/**
+ * @internal
+ */
+const get = (fromObject: any, path: string): any => {
+  let cursor = fromObject;
+  const pathComponents = path.split(".");
+  for (const step of pathComponents) {
+    if (!cursor || typeof cursor !== "object") {
+      return undefined;
+    }
+    cursor = cursor[step];
+  }
+  return cursor;
+};