Merge pull request #24 from JupiterOne/core-1836-optimistic-locking

CORE-1836 Optimistic locking with version number

Author: benrj

CHANGELOG.md

@@ -13,3 +13,9 @@ and this project adheres to
 ### Added
 
 - Support `consistentRead` option on `get` API
+
+## 1.5.0 - 2021-10-27
+
+### Added
+
+- Support optimistic locking for `put`, `update` and `delete` APIs

README.md

@@ -137,6 +137,43 @@ const { total } = await myDocumentDao.decr(
 );
 ```
 
+**Optimistic Locking with Version Numbers**
+
+For callers who wish to enable an optimistic locking strategy there are two
+available toggles:
+
+1. Provide the attribute you wish to be used to store the version number. This
+   will enable optimistic locking on the following operations: `put`, `update`,
+   and `delete`.
+
+   Writes for documents that do not have a version number attribute will
+   initialize the version number to 1. All subsequent writes will need to
+   provide the current version number. If an out-of-date version number is
+   supplied, an error will be thrown.
+
+   Example of Dao constructed with optimistic locking enabled.
+
+   ```
+   const dao = new DynamoDbDao<DataModel, KeySchema>({
+     tableName,
+     documentClient,
+     optimisticLockingAttribute: 'version',
+   });
+   ```
+
+2. If you wish to ignore optimistic locking for a save operation, specify
+   `ignoreOptimisticLocking: true` in the options on your `put`, `update`, or
+   `delete`.
+
+NOTE: Optimistic locking is NOT supported for `batchWrite` or `batchPut`
+operations. Consuming those APIs for data models that do have optimistic locking
+enabled may clobber your version data and could produce undesirable effects for
+other callers.
+
+This was modeled after the
+[Java Dynamo client](https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/DynamoDBMapper.OptimisticLocking.html)
+implementation.
+
 ## Developing
 
 The test setup requires that [docker-compose]() be installed. To run the tests,

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@jupiterone/dynamodb-dao",
-  "version": "1.4.0",
+  "version": "1.5.0",
   "description": "DynamoDB Data Access Object (DAO) helper library",
   "main": "index.js",
   "types": "index.d.ts",

src/index.test.ts

@@ -559,6 +559,82 @@ test('#generateUpdateParams should generate both update and remove params for do
   }
 });
 
+test('#generateUpdateParams should increment the version number', () => {
+  {
+    const options = {
+      tableName: 'blah2',
+      key: {
+        HashKey: 'abc',
+      },
+      data: {
+        a: 123,
+        b: 'abc',
+        c: undefined,
+        lockVersion: 1,
+      },
+      optimisticLockVersionAttribute: 'lockVersion',
+    };
+
+    expect(generateUpdateParams(options)).toEqual({
+      TableName: options.tableName,
+      Key: options.key,
+      ReturnValues: 'ALL_NEW',
+      ConditionExpression: '#lockVersion = :lockVersion',
+      UpdateExpression:
+        'add #lockVersion :lockVersionInc set #a0 = :a0, #a1 = :a1 remove #a2',
+      ExpressionAttributeNames: {
+        '#a0': 'a',
+        '#a1': 'b',
+        '#a2': 'c',
+        '#lockVersion': 'lockVersion',
+      },
+      ExpressionAttributeValues: {
+        ':a0': options.data.a,
+        ':a1': options.data.b,
+        ':lockVersionInc': 1,
+        ':lockVersion': 1,
+      },
+    });
+  }
+});
+
+test('#generateUpdateParams should increment the version number even when not supplied', () => {
+  {
+    const options = {
+      tableName: 'blah3',
+      key: {
+        HashKey: 'abc',
+      },
+      data: {
+        a: 123,
+        b: 'abc',
+        c: undefined,
+      },
+      optimisticLockVersionAttribute: 'lockVersion',
+    };
+
+    expect(generateUpdateParams(options)).toEqual({
+      TableName: options.tableName,
+      Key: options.key,
+      ReturnValues: 'ALL_NEW',
+      UpdateExpression:
+        'add #lockVersion :lockVersionInc set #a0 = :a0, #a1 = :a1 remove #a2',
+      ConditionExpression: 'attribute_not_exists(lockVersion)',
+      ExpressionAttributeNames: {
+        '#a0': 'a',
+        '#a1': 'b',
+        '#a2': 'c',
+        '#lockVersion': 'lockVersion',
+      },
+      ExpressionAttributeValues: {
+        ':a0': options.data.a,
+        ':a1': options.data.b,
+        ':lockVersionInc': 1,
+      },
+    });
+  }
+});
+
 test(`#queryUntilLimitReached should call #query if "filterExpression" not provided`, async () => {
   const keyConditionExpression = 'id = :id';
   const attributeValues = { id: uuid() };

src/index.ts

@@ -120,9 +120,57 @@ export interface ConditionalOptions {
   attributeValues?: AttributeValues;
 }
 
-export type PutOptions = ConditionalOptions;
-export type UpdateOptions = ConditionalOptions;
-export type DeleteOptions = ConditionalOptions;
+export interface SaveBehavior {
+  optimisticLockVersionAttribute?: string;
+  optimisticLockVersionIncrement?: number;
+}
+
+export interface MutateBehavior {
+  ignoreOptimisticLocking?: boolean;
+}
+
+export type PutOptions = ConditionalOptions & MutateBehavior;
+export type UpdateOptions = ConditionalOptions & MutateBehavior;
+export type DeleteOptions = ConditionalOptions & MutateBehavior;
+
+export interface BuildOptimisticLockOptionsInput extends ConditionalOptions {
+  versionAttribute: string;
+  versionAttributeValue: any;
+}
+
+export function buildOptimisticLockOptions(
+  options: BuildOptimisticLockOptionsInput
+): ConditionalOptions {
+  const { versionAttribute, versionAttributeValue } = options;
+  let { conditionExpression, attributeNames, attributeValues } = options;
+
+  const lockExpression = versionAttributeValue
+    ? `#${versionAttribute} = :${versionAttribute}`
+    : `attribute_not_exists(${versionAttribute})`;
+
+  conditionExpression = conditionExpression
+    ? `(${conditionExpression}) AND ${lockExpression}`
+    : lockExpression;
+
+  if (versionAttributeValue) {
+    attributeNames = {
+      ...attributeNames,
+      [`#${versionAttribute}`]: versionAttribute,
+    };
+    attributeValues = {
+      ...attributeValues,
+      [`:${versionAttribute}`]: versionAttributeValue,
+    };
+  }
+
+  return {
+    conditionExpression,
+    attributeNames,
+    attributeValues,
+  };
+}
+
+type DataModelAsMap = { [key: string]: any };
 
 export interface GenerateUpdateParamsInput extends UpdateOptions {
   tableName: string;
@@ -131,9 +179,10 @@ export interface GenerateUpdateParamsInput extends UpdateOptions {
 }
 
 export function generateUpdateParams(
-  options: GenerateUpdateParamsInput
+  options: GenerateUpdateParamsInput & SaveBehavior
 ): DocumentClient.UpdateItemInput {
   const setExpressions: string[] = [];
+  const addExpressions: string[] = [];
   const removeExpressions: string[] = [];
   const expressionAttributeNameMap: AttributeNames = {};
   const expressionAttributeValueMap: AttributeValues = {};
@@ -142,15 +191,41 @@ export function generateUpdateParams(
     tableName,
     key,
     data,
-    conditionExpression,
     attributeNames,
     attributeValues,
+    optimisticLockVersionAttribute: versionAttribute,
+    optimisticLockVersionIncrement: versionInc,
+    ignoreOptimisticLocking: ignoreLocking = false,
   } = options;
 
+  let conditionExpression = options.conditionExpression;
+
+  if (versionAttribute) {
+    addExpressions.push(`#${versionAttribute} :${versionAttribute}Inc`);
+    expressionAttributeNameMap[`#${versionAttribute}`] = versionAttribute;
+    expressionAttributeValueMap[`:${versionAttribute}Inc`] = versionInc ?? 1;
+
+    if (!ignoreLocking) {
+      ({ conditionExpression } = buildOptimisticLockOptions({
+        versionAttribute,
+        versionAttributeValue: (data as DataModelAsMap)[versionAttribute],
+        conditionExpression,
+      }));
+      expressionAttributeValueMap[`:${versionAttribute}`] = (
+        data as DataModelAsMap
+      )[versionAttribute];
+    }
+  }
+
   const keys = Object.keys(options.data).sort();
 
   for (let i = 0; i < keys.length; i++) {
     const name = keys[i];
+    if (name === versionAttribute) {
+      // versionAttribute is a special case and should always be handled
+      // explicitly as above with the supplied value ignored
+      continue;
+    }
 
     const valueName = `:a${i}`;
     const attributeName = `#a${i}`;
@@ -178,11 +253,13 @@ export function generateUpdateParams(
       ? 'remove ' + removeExpressions.join(', ')
       : undefined;
 
+  const addString =
+    addExpressions.length > 0 ? 'add ' + addExpressions.join(', ') : undefined;
   return {
     TableName: tableName,
     Key: key,
     ConditionExpression: conditionExpression,
-    UpdateExpression: [setString, removeString]
+    UpdateExpression: [addString, setString, removeString]
       .filter((val) => val !== undefined)
       .join(' '),
     ExpressionAttributeNames: {
@@ -197,9 +274,10 @@ export function generateUpdateParams(
   };
 }
 
-interface DynamoDbDaoInput {
+export interface DynamoDbDaoInput<T> {
   tableName: string;
   documentClient: DocumentClient;
+  optimisticLockingAttribute?: keyof NumberPropertiesInType<T>;
 }
 
 function invalidCursorError(cursor: string): Error {
@@ -261,10 +339,12 @@ export type NumberPropertiesInType<T> = Pick<
 export default class DynamoDbDao<DataModel, KeySchema> {
   public readonly tableName: string;
   public readonly documentClient: DocumentClient;
+  public readonly optimisticLockingAttribute?: keyof NumberPropertiesInType<DataModel>;
 
-  constructor(options: DynamoDbDaoInput) {
+  constructor(options: DynamoDbDaoInput<DataModel>) {
     this.tableName = options.tableName;
     this.documentClient = options.documentClient;
+    this.optimisticLockingAttribute = options.optimisticLockingAttribute;
   }
 
   /**
@@ -292,16 +372,30 @@ export default class DynamoDbDao<DataModel, KeySchema> {
    */
   async delete(
     key: KeySchema,
-    options: DeleteOptions = {}
+    options: DeleteOptions = {},
+    data: Partial<DataModel> = {}
   ): Promise<DataModel | undefined> {
+    let { attributeNames, attributeValues, conditionExpression } = options;
+
+    if (this.optimisticLockingAttribute && !options.ignoreOptimisticLocking) {
+      const versionAttribute = this.optimisticLockingAttribute.toString();
+      ({ attributeNames, attributeValues, conditionExpression } =
+        buildOptimisticLockOptions({
+          versionAttribute,
+          versionAttributeValue: (data as DataModelAsMap)[versionAttribute],
+          conditionExpression: conditionExpression,
+          attributeNames,
+          attributeValues,
+        }));
+    }
     const { Attributes: attributes } = await this.documentClient
       .delete({
         TableName: this.tableName,
         Key: key,
         ReturnValues: 'ALL_OLD',
-        ConditionExpression: options.conditionExpression,
-        ExpressionAttributeNames: options.attributeNames,
-        ExpressionAttributeValues: options.attributeValues,
+        ConditionExpression: conditionExpression,
+        ExpressionAttributeNames: attributeNames,
+        ExpressionAttributeValues: attributeValues,
       })
       .promise();
 
@@ -312,13 +406,36 @@ export default class DynamoDbDao<DataModel, KeySchema> {
    * Creates/Updates an item in the table
    */
   async put(data: DataModel, options: PutOptions = {}): Promise<DataModel> {
+    let { conditionExpression, attributeNames, attributeValues } = options;
+    if (this.optimisticLockingAttribute) {
+      // Must cast data to avoid tripping the linter, otherwise, it'll complain
+      // about expression of type 'string' can't be used to index type 'unknown'
+      const dataAsMap = data as DataModelAsMap;
+      const versionAttribute = this.optimisticLockingAttribute.toString();
+
+      if (!options.ignoreOptimisticLocking) {
+        ({ conditionExpression, attributeNames, attributeValues } =
+          buildOptimisticLockOptions({
+            versionAttribute,
+            versionAttributeValue: dataAsMap[versionAttribute],
+            conditionExpression,
+            attributeNames,
+            attributeValues,
+          }));
+      }
+
+      dataAsMap[versionAttribute] = dataAsMap[versionAttribute]
+        ? dataAsMap[versionAttribute] + 1
+        : 1;
+    }
+
     await this.documentClient
       .put({
         TableName: this.tableName,
         Item: data,
-        ConditionExpression: options.conditionExpression,
-        ExpressionAttributeNames: options.attributeNames,
-        ExpressionAttributeValues: options.attributeValues,
+        ConditionExpression: conditionExpression,
+        ExpressionAttributeNames: attributeNames,
+        ExpressionAttributeValues: attributeValues,
       })
       .promise();
     return data;
@@ -337,6 +454,8 @@ export default class DynamoDbDao<DataModel, KeySchema> {
       key,
       data,
       ...updateOptions,
+      optimisticLockVersionAttribute:
+        this.optimisticLockingAttribute?.toString(),
     });
     const { Attributes: attributes } = await this.documentClient
       .update(params)