Merge pull request #18 from lifeomic/batch-get

src/dynamo-table.test.ts

@@ -121,6 +121,75 @@ describe('DynamoTable', () => {
     });
   });
 
+  describe('batchGet', () => {
+    it('can fetch multiple items through batch api', async () => {
+      const batchGetSpy = jest.spyOn(dynamo.documentClient, 'batchGet');
+      const { userTable } = setupDb();
+
+      const ids = Array.from({ length: 101 }, (_, index) => String(index));
+
+      // seed the db
+      await userTable.batchPut(
+        ids.map((id) => ({
+          id,
+          account: 'account-1',
+          createdAt: new Date().toISOString(),
+        })),
+      );
+
+      const [first] = ids;
+
+      const batchGetFirst = await userTable.batchGet([{ id: first }], {
+        consistentRead: true,
+      });
+
+      expect(batchGetFirst).toStrictEqual([
+        { id: first, account: 'account-1', createdAt: expect.any(String) },
+      ]);
+      expect(batchGetSpy).toHaveBeenCalledTimes(1);
+      expect(batchGetSpy).toHaveBeenCalledWith({
+        RequestItems: {
+          'dynamost-user-table': {
+            ConsistentRead: true,
+            Keys: [
+              {
+                id: '0',
+              },
+            ],
+          },
+        },
+      });
+
+      batchGetSpy.mockClear();
+      const batchGetUnknown = await userTable.batchGet([{ id: 'unknown' }]);
+
+      expect(batchGetSpy).toHaveBeenCalledTimes(1);
+      expect(batchGetSpy).toHaveBeenCalledWith({
+        RequestItems: {
+          'dynamost-user-table': {
+            ConsistentRead: undefined,
+            Keys: [
+              {
+                id: 'unknown',
+              },
+            ],
+          },
+        },
+      });
+
+      expect(batchGetUnknown).toStrictEqual([]);
+
+      batchGetSpy.mockClear();
+      const batchGetOverLimit = await userTable.batchGet(
+        ids.map((id) => ({ id })),
+      );
+
+      expect(batchGetOverLimit.length).toBeGreaterThan(100);
+      expect(batchGetOverLimit).toHaveLength(ids.length);
+      expect(batchGetSpy).toHaveBeenCalledTimes(2);
+    });
+  });
+
   describe('upsert', () => {
     // TODO
   });

src/dynamo-table.ts

@@ -14,6 +14,8 @@ import {
 } from './dynamo-expressions';
 import { batchWrite } from './batch-write';
 import { Transaction } from './transaction-manager';
+import pMap from 'p-map';
+import chunk from 'lodash/chunk';
 
 /* -- Utility types to support indexing + other top-level config --  */
 export type KeySchema<Entity> = {
@@ -47,6 +49,10 @@ export type GetOptions = {
   consistentRead?: boolean;
 };
 
+export type BatchGetOptions = {
+  consistentRead?: boolean;
+};
+
 type AbstractZodOBject = z.ZodObject<any, any, any>;
 
 type BaseWriteOptions<Item> = {
@@ -340,6 +346,35 @@ export class DynamoTable<
     return this.schema.parse(result.Item);
   }
 
+  async batchGet(
+    keys: TableKey<this>[],
+    options?: BatchGetOptions,
+  ): Promise<z.infer<Schema>[]> {
+    // Only process 100 items at a time to avoid error:
+    // Too many items requested for the BatchWriteItem call
+    const batchGetMaxItemCount = 100;
+
+    const results = await pMap(
+      chunk(keys, batchGetMaxItemCount),
+      (requests) =>
+        this.client.batchGet({
+          RequestItems: {
+            [this.config.tableName]: {
+              ConsistentRead: options?.consistentRead,
+              Keys: requests,
+            },
+          },
+        }),
+      {
+        concurrency: 2,
+      },
+    );
+
+    return results
+      .flatMap((set) => set.Responses?.[this.config.tableName] || [])
+      .map((item) => this.schema.parse(item));
+  }
+
   private async _query(params: {
     index?: string;
     key: KeyCondition<z.infer<Schema>, any>;