Merge pull request #88 from shiftcode/#83-improve-batch-requests

#83 improve batch requests

Author: michaelwittwer

src/dynamo/batchget/batch-get-full.response.ts

@@ -6,7 +6,7 @@ export interface BatchGetFullResponse {
   /**
    * A map of table name to a list of items. Each object in Responses consists of a table name, along with a map of attribute data consisting of the data type and attribute value.
    */
-  Responses?: BatchGetResponse
+  Responses: BatchGetResponse
   /**
    * A map of tables and their respective keys that were not processed with the current response. The UnprocessedKeys value is in the same form as RequestItems, so the value can be provided directly to a subsequent BatchGetItem operation. For more information, see RequestItems in the Request Parameters section. Each element consists of:    Keys - An array of primary key attribute values that define specific items in the table.    ProjectionExpression - One or more attributes to be retrieved from the table or index. By default, all attributes are returned. If a requested attribute is not found, it does not appear in the result.    ConsistentRead - The consistency of a read operation. If set to true, then a strongly consistent read is used; otherwise, an eventually consistent read is used.   If there are no unprocessed keys remaining, the response contains an empty UnprocessedKeys map.
    */

src/dynamo/batchget/batch-get-utils.spec.ts

@@ -0,0 +1,118 @@
+import { DynamoDB } from 'aws-sdk'
+import { of } from 'rxjs'
+import { DynamoRx } from '../dynamo-rx'
+import { batchGetItemsFetchAll, combineBatchGetResponses, hasUnprocessedKeys } from './batch-get-utils'
+
+describe('batch-get utils', () => {
+
+  describe('hasUnprocessedKeys', () => {
+    it('should return bool according to given object', () => {
+      expect(hasUnprocessedKeys({})).toBeFalsy()
+      expect(hasUnprocessedKeys({ Responses: {} })).toBeFalsy()
+      expect(hasUnprocessedKeys({ UnprocessedKeys: {} })).toBeFalsy()
+      expect(hasUnprocessedKeys({ UnprocessedKeys: { 'aTableName': { Keys: [] } } })).toBeFalsy()
+      expect(hasUnprocessedKeys({ UnprocessedKeys: { 'aTableName': { Keys: [{ id: { S: 'id' } }] } } })).toBeTruthy()
+    })
+  })
+
+  describe('combineBatchGetResponses', () => {
+    const resp1: DynamoDB.BatchGetItemOutput = {
+      Responses: {
+        'tableA': [
+          { id: { S: 'id-a1' } },
+        ],
+        'tableB': [
+          { id: { S: 'id-b' } },
+        ],
+      },
+      UnprocessedKeys: {
+        'tableA:': { Keys: [{ id: { S: 'id-a2' } }] },
+        'tableC:': { Keys: [{ id: { S: 'id-c' } }] },
+        'tableD:': { Keys: [{ id: { S: 'id-d' } }] },
+      },
+    }
+    const resp2: DynamoDB.BatchGetItemOutput = {
+      Responses: {
+        'tableA': [
+          { id: { S: 'id-a2' } },
+        ],
+        'tableC': [
+          { id: { S: 'id-c' } },
+        ],
+      },
+      UnprocessedKeys: {
+        'tableD:': { Keys: [{ id: { S: 'id-d' } }] },
+      },
+    }
+    const expectedOutput: DynamoDB.BatchGetItemOutput = {
+      Responses: {
+        'tableA': [
+          { id: { S: 'id-a1' } },
+          { id: { S: 'id-a2' } },
+        ],
+        'tableB': [
+          { id: { S: 'id-b' } },
+        ],
+        'tableC': [
+          { id: { S: 'id-c' } },
+        ],
+      },
+      UnprocessedKeys: {
+        'tableD:': { Keys: [{ id: { S: 'id-d' } }] },
+      },
+    }
+    it('should combine correctly', () => {
+      expect(combineBatchGetResponses(resp1)(resp2)).toEqual(expectedOutput)
+    })
+
+  })
+
+  describe('batchGetItemsFetchAll', () => {
+    let batchGetItemsSpy: jasmine.Spy
+    let dynamoRx: DynamoRx
+    let backoffTimerMock: { next: jasmine.Spy }
+
+    const output1: DynamoDB.BatchGetItemOutput = {
+      Responses: {
+        'tableA': [{ id: { S: 'id-A' } }],
+      },
+      UnprocessedKeys: {
+        'tableA': {
+          Keys: [{ id: { S: 'id-A' } }],
+        },
+      },
+    }
+    const output2: DynamoDB.BatchGetItemOutput = {
+      Responses: {
+        'tableA': [{ id: { S: 'id-A' } }],
+      },
+    }
+
+    beforeEach(async () => {
+      batchGetItemsSpy = jasmine.createSpy().and.returnValues(of(output1), of(output2))
+      dynamoRx = <any>{ batchGetItems: batchGetItemsSpy }
+      backoffTimerMock = { next: jasmine.createSpy().and.returnValue({ value: 0 }) }
+
+      await batchGetItemsFetchAll(
+        dynamoRx,
+        <any>{},
+        <IterableIterator<number>><any>backoffTimerMock,
+        0,
+      ).toPromise()
+    })
+
+
+    it('should use UnprocessedKeys for next request', () => {
+      expect(batchGetItemsSpy).toHaveBeenCalledTimes(2)
+      expect(batchGetItemsSpy.calls.mostRecent().args[0]).toBeDefined()
+      expect(batchGetItemsSpy.calls.mostRecent().args[0].RequestItems).toBeDefined()
+      expect(batchGetItemsSpy.calls.mostRecent().args[0].RequestItems).toEqual(output1.UnprocessedKeys)
+    })
+
+    it('should backoff when UnprocessedItems', () => {
+      expect(backoffTimerMock.next).toHaveBeenCalledTimes(1)
+    })
+
+  })
+
+})

src/dynamo/batchget/batch-get-utils.ts

@@ -0,0 +1,67 @@
+import { DynamoDB } from 'aws-sdk'
+import { BatchGetRequestMap } from 'aws-sdk/clients/dynamodb'
+import { Observable, of } from 'rxjs'
+import { delay, map, mergeMap } from 'rxjs/operators'
+import { DynamoRx } from '../dynamo-rx'
+
+export function batchGetItemsFetchAll(
+  dynamoRx: DynamoRx,
+  params: DynamoDB.BatchGetItemInput,
+  backoffTimer: IterableIterator<number>,
+  throttleTimeSlot: number,
+): Observable<DynamoDB.BatchGetItemOutput> {
+  return dynamoRx.batchGetItems(params)
+    .pipe(
+      mergeMap(response => {
+        if (hasUnprocessedKeys(response)) {
+          return of(response.UnprocessedKeys)
+            .pipe(
+              delay(backoffTimer.next().value * throttleTimeSlot),
+              mergeMap((UnprocessedKeys: DynamoDB.BatchGetRequestMap) => {
+                const nextParams = { ...params, RequestItems: UnprocessedKeys }
+                return batchGetItemsFetchAll(dynamoRx, nextParams, backoffTimer, throttleTimeSlot)
+              }),
+              map(combineBatchGetResponses(response)),
+            )
+        }
+        return of(response)
+      }),
+    )
+}
+
+export type ResponseWithUnprocessedKeys = DynamoDB.BatchGetItemOutput & { UnprocessedKeys: BatchGetRequestMap }
+
+export function hasUnprocessedKeys(response: DynamoDB.BatchGetItemOutput): response is ResponseWithUnprocessedKeys {
+  if (!response.UnprocessedKeys) {
+    return false
+  }
+  return Object.values(response.UnprocessedKeys)
+    .some(t => !!t && t.Keys && t.Keys.length > 0)
+}
+
+/**
+ * combines a first with a second response. ConsumedCapacity is always from the latter.
+ * @param response1
+ */
+export function combineBatchGetResponses(response1: DynamoDB.BatchGetItemOutput) {
+  return (response2: DynamoDB.BatchGetItemOutput): DynamoDB.BatchGetItemOutput => {
+    const tableNames: string[] = Object.keys(response1.Responses || {})
+
+    Object.keys(response2.Responses || {})
+      .filter(tn => !tableNames.includes(tn))
+      .forEach(tn => tableNames.push(tn))
+
+    const Responses = tableNames
+      .reduce((u, tableName) => ({
+        ...u,
+        [tableName]: [
+          ...(response1.Responses && response1.Responses[tableName] || []),
+          ...(response2.Responses && response2.Responses[tableName] || []),
+        ],
+      }), {})
+    return {
+      ...response2,
+      Responses,
+    }
+  }
+}

src/dynamo/batchget/batch-get.request.spec.ts

@@ -1,16 +1,15 @@
+// tslint:disable:no-non-null-assertion
 import * as DynamoDB from 'aws-sdk/clients/dynamodb'
 import { of } from 'rxjs'
 import { getTableName } from '../../../test/helper'
-import { SimpleWithCompositePartitionKeyModel, SimpleWithPartitionKeyModel } from '../../../test/models'
-import { Organization } from '../../../test/models/organization.model'
-import { Attributes } from '../../mapper'
+import { Organization, SimpleWithCompositePartitionKeyModel, SimpleWithPartitionKeyModel } from '../../../test/models'
+import { Attributes, toDb } from '../../mapper'
 import { DynamoRx } from '../dynamo-rx'
 import { BatchGetRequest } from './batch-get.request'
 
 describe('batch get', () => {
   let request: BatchGetRequest
 
-
   describe('params', () => {
 
     beforeEach(() => request = new BatchGetRequest())
@@ -21,39 +20,150 @@ describe('batch get', () => {
     })
 
     it('key', () => {
-      request.forModel(Organization, ['idValue'])
+      const o: Partial<Organization> = {
+        id: 'idValue',
+        createdAtDate: new Date(),
+      }
+      request.forModel(Organization, [o])
       const params = request.params
       expect(params.RequestItems).toBeDefined()
       expect(params.RequestItems.Organization).toBeDefined()
-      expect(params.RequestItems.Organization).toEqual({ Keys: [{ id: { S: 'idValue' } }] })
+      expect(params.RequestItems.Organization.Keys).toBeDefined()
+      expect(params.RequestItems.Organization.Keys).toEqual([{
+        id: { S: 'idValue' },
+        createdAtDate: { S: o.createdAtDate!.toISOString() },
+      }])
     })
   })
 
-
   describe('forModel', () => {
     beforeEach(() => request = new BatchGetRequest())
 
     it('should throw when same table is used 2 times', () => {
-      request.forModel(SimpleWithPartitionKeyModel, ['idVal'])
-      expect(() => request.forModel(SimpleWithPartitionKeyModel, ['otherVal'])).toThrow()
+      request.forModel(SimpleWithPartitionKeyModel, [{ id: 'idVal' }])
+      expect(() => request.forModel(SimpleWithPartitionKeyModel, [{ id: 'otherVal' }])).toThrow()
+    })
+
+    it('should throw when sortKey is missing but necessary', () => {
+      expect(() => request.forModel(SimpleWithCompositePartitionKeyModel, [{ id: 'idVal' }]))
+    })
+
+    it('should throw when modelClazz is not @Model decorated', () => {
+      class X {id: string}
+
+      expect(() => request.forModel(X, [{ id: 'ok' }])).toThrow()
     })
 
     it('should throw when providing null value ', () => {
       expect(() => request.forModel(SimpleWithPartitionKeyModel, [<any>null])).toThrow()
     })
 
-    it('should throw when sortKey is missing', () => {
-      expect(() => request.forModel(SimpleWithCompositePartitionKeyModel, [{ partitionKey: 'idVal' }]))
+    it('should allow ConsistentRead', () => {
+      request.forModel(SimpleWithPartitionKeyModel, [{ id: 'myId' }], true)
+      expect(request.params).toBeDefined()
+      expect(request.params.RequestItems).toBeDefined()
+      const keysOfTable = request.params.RequestItems[getTableName(SimpleWithPartitionKeyModel)]
+      expect(keysOfTable).toBeDefined()
+      expect(keysOfTable.ConsistentRead).toBeTruthy()
     })
 
-    it('should throw when partitionKey is neither string nor object', () => {
-      expect(() => request.forModel(SimpleWithCompositePartitionKeyModel, [<any>78]))
-      expect(() => request.forModel(SimpleWithCompositePartitionKeyModel, [<any>true]))
-      expect(() => request.forModel(SimpleWithCompositePartitionKeyModel, [<any>new Date()]))
+    it('should throw when more than 100 items are added', () => {
+      const items55: Array<Partial<SimpleWithPartitionKeyModel>> = new Array(55)
+        .map((x, i) => ({ id: `id-${i}` }))
+      const items60: Array<Partial<Organization>> = new Array(60)
+        .map((x, i) => ({ id: `id-${i}`, createdAtDate: new Date() }))
+
+      // at once
+      expect(() => request.forModel(SimpleWithPartitionKeyModel, [...items55, ...items55])).toThrow()
+
+      // in two steps
+      expect(() => {
+        request.forModel(SimpleWithPartitionKeyModel, items55)
+        request.forModel(Organization, items60)
+      }).toThrow()
     })
 
   })
 
+  describe('execNoMap, execFullResponse, exec', () => {
+    const jsItem1: SimpleWithPartitionKeyModel = { id: 'id-1', age: 21 }
+    const jsItem2: SimpleWithPartitionKeyModel = { id: 'id-2', age: 22 }
+
+    const output1: DynamoDB.BatchGetItemOutput = {
+      Responses: {
+        [getTableName(SimpleWithPartitionKeyModel)]: [toDb(jsItem1, SimpleWithPartitionKeyModel)],
+      },
+      UnprocessedKeys: {
+        [getTableName(SimpleWithPartitionKeyModel)]: {
+          Keys: [toDb(jsItem1, SimpleWithPartitionKeyModel)],
+        },
+      },
+    }
+    const output2: DynamoDB.BatchGetItemOutput = {
+      Responses: {
+        [getTableName(SimpleWithPartitionKeyModel)]: [toDb(jsItem2, SimpleWithPartitionKeyModel)],
+      },
+    }
+
+    let batchGetItemsSpy: jasmine.Spy
+    let nextSpyFn: () => { value: number }
+
+    const generatorMock = () => <any>{ next: nextSpyFn }
+
+    beforeEach(() => {
+      request = new BatchGetRequest()
+      request.forModel(SimpleWithPartitionKeyModel, [jsItem1, jsItem2])
+
+      batchGetItemsSpy = jasmine.createSpy().and.returnValues(of(output1), of(output2))
+      const dynamoRx: DynamoRx = <any>{ batchGetItems: batchGetItemsSpy }
+
+      Object.assign(request, { dynamoRx })
+
+      nextSpyFn = jest.fn().mockImplementation(() => ({ value: 0 }))
+    })
+
+    it('[execNoMap] should backoff and retry when UnprocessedItems are returned', async () => {
+      const result = await request.execNoMap(generatorMock).toPromise()
+      expect(nextSpyFn).toHaveBeenCalledTimes(1)
+      expect(batchGetItemsSpy).toHaveBeenCalledTimes(2)
+      expect(result).toBeDefined()
+      expect(result.Responses).toBeDefined()
+
+      const resultItems = result.Responses![getTableName(SimpleWithPartitionKeyModel)]
+      expect(resultItems).toBeDefined()
+      expect(resultItems.length).toBe(2)
+      expect(resultItems[0]).toEqual(toDb(jsItem1, SimpleWithPartitionKeyModel))
+      expect(resultItems[1]).toEqual(toDb(jsItem2, SimpleWithPartitionKeyModel))
+    })
+
+    it('[execFullResponse] should backoff and retry when UnprocessedItems are returned', async () => {
+      const result = await request.execFullResponse(generatorMock).toPromise()
+      expect(nextSpyFn).toHaveBeenCalledTimes(1)
+      expect(batchGetItemsSpy).toHaveBeenCalledTimes(2)
+      expect(result).toBeDefined()
+      expect(result.Responses).toBeDefined()
+
+      const resultItems = result.Responses![getTableName(SimpleWithPartitionKeyModel)]
+      expect(resultItems).toBeDefined()
+      expect(resultItems.length).toBe(2)
+      expect(resultItems[0]).toEqual(jsItem1)
+      expect(resultItems[1]).toEqual(jsItem2)
+    })
+
+    it('[exec] should backoff and retry when UnprocessedItems are returned', async () => {
+      const result = await request.exec(generatorMock).toPromise()
+      expect(nextSpyFn).toHaveBeenCalledTimes(1)
+      expect(batchGetItemsSpy).toHaveBeenCalledTimes(2)
+      expect(result).toBeDefined()
+
+      const resultItems = result[getTableName(SimpleWithPartitionKeyModel)]
+      expect(resultItems).toBeDefined()
+      expect(resultItems.length).toBe(2)
+      expect(resultItems[0]).toEqual(jsItem1)
+      expect(resultItems[1]).toEqual(jsItem2)
+    })
+
+  })
 
   describe('should map the result items', () => {
     let batchGetItemsSpy: jasmine.Spy
@@ -74,7 +184,7 @@ describe('batch get', () => {
       const dynamoRx: DynamoRx = <any>{ batchGetItems: batchGetItemsSpy }
       request = new BatchGetRequest()
       Object.assign(request, { dynamoRx })
-      request.forModel(SimpleWithPartitionKeyModel, ['idVal'])
+      request.forModel(SimpleWithPartitionKeyModel, [{ id: 'idVal' }])
     })
 
     it('exec', async () => {
@@ -93,5 +203,4 @@ describe('batch get', () => {
     })
 
   })
-
 })