fix: add batching to clearAllItems (#629)

* fix: add batching to clearAllItems * fix: add types to chunk helper
erezrokah · Jun 22, 2022 · 54b37b7 · 54b37b7
1 parent dc2c8b8
commit 54b37b7
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 6 deletions.
diff --git a/src/utils/dynamoDb.test.ts b/src/utils/dynamoDb.test.ts
@@ -85,6 +85,52 @@ describe('dynamoDb utils', () => {
         },
       });
     });
+    test('should call batchWrite multiple times for db containing more than 25 items', async () => {
+      const describeTable = db().describeTable;
+      const describeTablePromise = describeTable().promise;
+      const table = { KeySchema: [{ AttributeName: 'id' }] };
+      describeTablePromise.mockReturnValue(Promise.resolve({ Table: table }));
+
+      const scan = documentClient().scan;
+      const batchWrite = documentClient().batchWrite;
+      const scanPromise = scan().promise;
+      const items = Array(30)
+        .fill({ id: 'id' })
+        .map((el, i) => ({
+          id: el.id + i,
+        }));
+      scanPromise.mockReturnValue(Promise.resolve({ Items: items }));
+
+      jest.clearAllMocks();
+
+      await clearAllItems(region, tableName);
+
+      expect(scan).toHaveBeenCalledTimes(1);
+      expect(scan).toHaveBeenCalledWith({
+        AttributesToGet: table.KeySchema.map((k) => k.AttributeName),
+        TableName: tableName,
+      });
+
+      expect(batchWrite).toHaveBeenCalledTimes(2);
+      const batch1 = items.splice(0, 25);
+      const batch2 = items;
+      expect(batch1.length).toBe(25);
+      expect(batch2.length).toBe(5);
+      expect(batchWrite).toHaveBeenCalledWith({
+        RequestItems: {
+          [tableName]: batch1.map((item) => ({
+            DeleteRequest: { Key: { id: item.id } },
+          })),
+        },
+      });
+      expect(batchWrite).toHaveBeenCalledWith({
+        RequestItems: {
+          [tableName]: batch2.map((item) => ({
+            DeleteRequest: { Key: { id: item.id } },
+          })),
+        },
+      });
+    });
   });
 
   describe('writeItems', () => {

diff --git a/src/utils/dynamoDb.ts b/src/utils/dynamoDb.ts
@@ -1,5 +1,11 @@
 import AWS = require('aws-sdk');
 
+function* chunks<T>(arr: T[], n: number): Generator<T[], void> {
+  for (let i = 0; i < arr.length; i += n) {
+    yield arr.slice(i, i + n);
+  }
+}
+
 const itemToKey = (
   item: AWS.DynamoDB.DocumentClient.AttributeMap,
   keySchema: AWS.DynamoDB.KeySchemaElement[],
@@ -31,13 +37,15 @@ export const clearAllItems = async (region: string, tableName: string) => {
   const items = scanResult.Items || [];
 
   if (items.length > 0) {
-    const deleteRequests = items.map((item) => ({
-      DeleteRequest: { Key: itemToKey(item, keySchema) },
-    }));
+    for (const chunk of chunks(items, 25)) {
+      const deleteRequests = chunk.map((item) => ({
+        DeleteRequest: { Key: itemToKey(item, keySchema) },
+      }));
 
-    await db
-      .batchWrite({ RequestItems: { [tableName]: deleteRequests } })
-      .promise();
+      await db
+        .batchWrite({ RequestItems: { [tableName]: deleteRequests } })
+        .promise();
+    }
   }
 };