Merge pull request #3 from emdgroup/fix/gsi-pagination

fix pagination on index
emdgroup · Apr 23, 2023 · 7e1906f · 7e1906f
2 parents 7cd7892 + bd8188c
commit 7e1906f
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 6 deletions.
diff --git 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
@@ -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
@@ -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;