From bd8188cd6ba263224eb84a54afb50424a096c21a Mon Sep 17 00:00:00 2001
From: Moritz Onken <onken@netcubed.de>
Date: Sun, 23 Apr 2023 15:37:20 -0400
Subject: [PATCH] fix pagination on index

---
 package.json      |  2 +-
 src/index.spec.ts | 16 ++++++++++++----
 src/index.ts      |  6 +++++-
 3 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/package.json b/package.json
index d1ca55f..667c4fa 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@emdgroup/dynamodb-paginator",
-  "version": "2.1.1",
+  "version": "2.1.2",
   "main": "dist/cjs/index.js",
   "exports": {
     "import": "./dist/esm/index.js",
diff --git a/src/index.spec.ts b/src/index.spec.ts
index 8d4149a..d671f88 100644
--- a/src/index.spec.ts
+++ b/src/index.spec.ts
@@ -147,11 +147,19 @@ describe('aws', () => {
             assert.equal(paginator.scannedCount, 20);
         });
 
-        it('nextToken', async () => {
+        [null, 'GSI1.1'].map((idx) => it(idx ? 'nextToken with index' : 'nextToken', async () => {
+            // give index some time to update
+            if (idx) await new Promise((resolve) => setTimeout(resolve, 1000));
             const paginator = ({ from, context }: { from?: string, context?: string } = {}) => paginateQuery({
                 TableName: TABLE_NAME,
-                KeyConditionExpression: 'PK = :pk',
-                FilterExpression: 'GSI1SK >= :sk',
+                ...idx ? {
+                    IndexName: 'GSI1.1',
+                    KeyConditionExpression: 'GSI1PK = :pk',
+                    FilterExpression: 'SK >= :sk',
+                } : {
+                    KeyConditionExpression: 'PK = :pk',
+                    FilterExpression: 'GSI1SK >= :sk',
+                },
                 ExpressionAttributeValues: {
                     ':pk': 'p:',
                     ':sk': items[4].SK,
@@ -184,7 +192,7 @@ describe('aws', () => {
 
             const pageCtx = paginator({ from: page1.nextToken, context: 'foobar' });
             await assert.rejects(() => pageCtx.all(), { name: 'TokenError' });
-        });
+        }).retries(10));
 
         it('predicate', async () => {
             type Foo = { PK: string; SK: string };
diff --git a/src/index.ts b/src/index.ts
index 7f7214a..bd2f313 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -343,7 +343,11 @@ export class PaginationResponse<T extends AttributeMap = AttributeMap> {
             while (items.length) {
                 const item = items.shift();
                 if (!item) continue;
-                this._nextKey = this.buildLastEvaluatedKey(item, query.IndexName);
+                this._nextKey = this.buildLastEvaluatedKey(item);
+                if (query.IndexName) this._nextKey = {
+                    ...this._nextKey,
+                    ...this.buildLastEvaluatedKey(item, query.IndexName),
+                };
                 if (filter && !filter(item)) continue;
                 this.count += 1;
                 yield item as T;