feat: add support for using add and remove update operators with sets

src/decorator/attribute-types/string-set.spec.ts

@@ -9,6 +9,12 @@ describe('AttributeType/StringSet', () => {
     record = TestableTable.new()
   })
 
+  afterEach(async () => {
+    if (!record.isNew()) {
+      await record.delete()
+    }
+  })
+
   describe('setting', () => {
     it('should allow values to be a Set', () => {
       expect(record.testStringSet).eq(null)
@@ -36,4 +42,28 @@ describe('AttributeType/StringSet', () => {
       expect(isArray(record.testStringSetArray)).eq(true)
     })
   })
+
+  describe('update operators', () => {
+    it('should support update operators add', async () => {
+      record.id = 10
+      record.title = 'add to set'
+      record.generic = 'something'
+      expect(record.testStringSet).eq(null)
+      record.testStringSet = new Set<string>(['some value'])
+      await record.save()
+
+      record.set('testStringSet', record.testStringSet.add('some new value'), {
+        operator: 'add',
+      })
+
+      const updateInput = TestableTable.documentClient.getUpdateInput(record)
+      expect(updateInput.UpdateExpression).eq('SET #UA1 = :u1 ADD #UA0 :u0')
+      expect(updateInput.ExpressionAttributeValues![':u0'].SS).deep.eq(['some value', 'some new value'])
+
+      await record.save()
+
+      const reloaded = await TestableTable.primaryKey.get({ id: 10, title: 'add to set' })
+      expect(toArray(reloaded?.testStringSet)).to.deep.eq(['some new value', 'some value'])
+    })
+  })
 })

src/interfaces/update-operator.interface.ts

@@ -8,5 +8,9 @@
  * For List attributes:
  * * ``append`` - Append one or more values to a List attribute
  * * ``if_not_exists`` - Increment the stored value by the specified amount
+ *
+ * For Set attributes:
+ * * ``add`` - Add one or more values from a Set attribute
+ * * ``delete`` - Delete one or more values from a Set attribute
  */
-export type UpdateOperator = 'append' | 'decrement' | 'if_not_exists' | 'increment' | 'set'
+export type UpdateOperator = 'append' | 'add' | 'delete' | 'decrement' | 'if_not_exists' | 'increment' | 'set'

src/query/update-item-input.ts

@@ -28,6 +28,8 @@ export function getUpdateItemInput<T extends Table>(record: T, params?: UpdateIt
   }
 
   const sets: string[] = []
+  const adds: string[] = []
+  const deletes: string[] = []
   const removes: string[] = []
   const attributeNameMap: Record<string, string> = {}
   const attributeValueMap: AttributeMap = {}
@@ -40,7 +42,7 @@ export function getUpdateItemInput<T extends Table>(record: T, params?: UpdateIt
   _.each(_.uniq(record.getUpdatedAttributes()), (attributeName, i) => {
     const attribute = tableClass.schema.getAttributeByName(attributeName)
     const value = attribute.toDynamo(record.getAttribute(attributeName))
-    const operator = record.getUpdateOperator(attributeName)
+    const operator = record.getAttributeUpdateOperator(attributeName)
     const slug = `#UA${valueCounter}`
 
     if (value != null) {
@@ -56,6 +58,10 @@ export function getUpdateItemInput<T extends Table>(record: T, params?: UpdateIt
         case 'append': sets.push(`${slug} = list_append(${slug}, :u${valueCounter})`); break
         case 'if_not_exists': sets.push(`${slug} = if_not_exists(${slug}, :u${valueCounter})`); break
 
+        // Set attribute operators
+        case 'add': adds.push(`${slug} :u${valueCounter}`); break
+        case 'delete': deletes.push(`${slug} :u${valueCounter}`); break
+
         case 'set':
         default: sets.push(`${slug} = :u${valueCounter}`); break
       }
@@ -77,6 +83,22 @@ export function getUpdateItemInput<T extends Table>(record: T, params?: UpdateIt
     updateExpression += 'SET ' + sets.join(', ')
   }
 
+  if (adds.length > 0) {
+    if (updateExpression.length > 0) {
+      updateExpression += ' '
+    }
+
+    updateExpression += 'ADD ' + adds.join(', ')
+  }
+
+  if (removes.length > 0) {
+    if (updateExpression.length > 0) {
+      updateExpression += ' '
+    }
+
+    updateExpression += 'REMOVE ' + removes.join(', ')
+  }
+
   if (removes.length > 0) {
     if (updateExpression.length > 0) {
       updateExpression += ' '

src/table.ts

@@ -352,10 +352,23 @@ export class Table {
   /**
    * Get the update operator for an attribute.
    */
-  public getUpdateOperator(attributeName: string): UpdateOperator {
+  public getUpdateOperator<P extends TableProperty<this>>(propertyName: P | string): UpdateOperator {
+    const attribute = this.table.schema.getAttributeByPropertyName(propertyName as string)
+    return this.getAttributeUpdateOperator(attribute.name)
+  }
+
+  public getAttributeUpdateOperator(attributeName: string): UpdateOperator {
     return this.__updateOperators[attributeName] ?? 'set'
   }
 
+  /**
+   * Set the update operator for a property.
+   */
+  public setUpdateOperator<P extends TableProperty<this>>(propertyName: P | string, operator: UpdateOperator): this {
+    const attribute = this.table.schema.getAttributeByPropertyName(propertyName as string)
+    return this.setAttributeUpdateOperator(attribute.name, operator)
+  }
+
   /**
    * Set the update operator for an attribute.
    */
@@ -564,11 +577,16 @@ export class Table {
       let output: PutItemCommandOutput | UpdateItemCommandOutput
       if (beforeSaveEvent.operator === 'put') {
         output = await this.table.documentClient.put(this, beforeSaveEvent)
-        this.__putRequired = false
       } else {
         output = await this.table.documentClient.update(this, beforeSaveEvent)
       }
 
+      // reset internal tracking of changes attributes
+      this.__putRequired = false
+      this.__removedAttributes = []
+      this.__updatedAttributes = []
+      this.__updateOperators = {}
+
       // trigger afterSave before clearing values, so the hook can determine what has been changed
       await this.afterSave({
         ...beforeSaveEvent,
@@ -577,11 +595,6 @@ export class Table {
         updatedAttributes: this.__updatedAttributes,
       })
 
-      // reset internal tracking of changes attributes
-      this.__removedAttributes = []
-      this.__updatedAttributes = []
-      this.__updateOperators = {}
-
       if (beforeSaveEvent.returnOutput === true) {
         return output
       }