Support add/set/delete operations to update item in transaction (#718)

README_transact.md

@@ -49,6 +49,32 @@ Dynamoid::TransactionWrite.execute do |txn|
 end
 ```
 
+### Update items
+An item can be updated by providing the hash key, range key if applicable, and the fields to update.
+Updating fields can also be done within a block using the `set()` method.
+To increment a numeric value or to add values to a set use `add()` within the block.
+Similarly a field can be removed or values can be removed from a set by using `delete()` in the block.
+```ruby
+Dynamoid::TransactionWrite.execute do |txn|
+  # sets the name and title for user 1
+  # The user is found by id
+  txn.update!(User, id: 1, name: 'bob', title: 'mister')
+
+  # sets the name, increments a count and deletes a field
+  txn.update!(user) do |u| # a User instance is provided
+    u.set(name: 'bob')
+    u.add(article_count: 1)
+    u.delete(:title)
+  end
+
+  # adds to a set of integers and deletes from a set of strings
+  txn.update!(User, id: 3) do |u|
+    u.add(friend_ids: [1, 2])
+    u.delete(child_names: ['bebe'])
+  end
+end
+```
+
 ### Destroy or delete items
 Models can be used or the model class and key can be specified.
 When the key is a single column it is specified as a single value or a hash

lib/dynamoid/transaction_write.rb

@@ -38,26 +38,26 @@ def save(model, options = {})
       model.new_record? ? create(model, {}, options) : update(model, {}, options)
     end
 
-    def create!(model_or_model_class, attributes = {}, options = {})
-      create(model_or_model_class, attributes, options.reverse_merge(raise_validation_error: true))
+    def create!(model_or_model_class, attributes = {}, options = {}, &block)
+      create(model_or_model_class, attributes, options.reverse_merge(raise_validation_error: true), &block)
     end
 
-    def create(model_or_model_class, attributes = {}, options = {})
-      add_action_and_validate Dynamoid::TransactionWrite::Create.new(model_or_model_class, attributes, options)
+    def create(model_or_model_class, attributes = {}, options = {}, &block)
+      add_action_and_validate Dynamoid::TransactionWrite::Create.new(model_or_model_class, attributes, options, &block)
     end
 
     # upsert! does not exist because upserting instances that can raise validation errors is not officially supported
 
-    def upsert(model_or_model_class, attributes = {}, options = {})
-      add_action_and_validate Dynamoid::TransactionWrite::Upsert.new(model_or_model_class, attributes, options)
+    def upsert(model_or_model_class, attributes = {}, options = {}, &block)
+      add_action_and_validate Dynamoid::TransactionWrite::Upsert.new(model_or_model_class, attributes, options, &block)
     end
 
-    def update!(model_or_model_class, attributes = {}, options = {})
-      update(model_or_model_class, attributes, options.reverse_merge(raise_validation_error: true))
+    def update!(model_or_model_class, attributes = {}, options = {}, &block)
+      update(model_or_model_class, attributes, options.reverse_merge(raise_validation_error: true), &block)
     end
 
-    def update(model_or_model_class, attributes = {}, options = {})
-      add_action_and_validate Dynamoid::TransactionWrite::Update.new(model_or_model_class, attributes, options)
+    def update(model_or_model_class, attributes = {}, options = {}, &block)
+      add_action_and_validate Dynamoid::TransactionWrite::Update.new(model_or_model_class, attributes, options, &block)
     end
 
     def delete(model_or_model_class, key_or_attributes = {}, options = {})

lib/dynamoid/transaction_write/action.rb

@@ -4,7 +4,7 @@ module Dynamoid
   class TransactionWrite
     class Action
       VALID_OPTIONS = %i[skip_callbacks skip_validation raise_validation_error skip_existence_check].freeze
-      attr_accessor :model, :attributes, :options
+      attr_accessor :model, :attributes, :options, :additions, :deletions, :removals
 
       def initialize(model_or_model_class, attributes = {}, options = {})
         if model_or_model_class.is_a?(Dynamoid::Document)
@@ -14,8 +14,13 @@ def initialize(model_or_model_class, attributes = {}, options = {})
         end
         self.attributes = attributes
         self.options = options || {}
+        self.additions = {}
+        self.deletions = {}
+        self.removals = []
         invalid_keys = self.options.keys - VALID_OPTIONS
         raise ArgumentError, "Invalid options found: '#{invalid_keys}'" if invalid_keys.present?
+
+        yield(self) if block_given?
       end
 
       def model_class
@@ -37,6 +42,26 @@ def range_key
         attributes[model_class.range_key]
       end
 
+      # sets a value in the attributes
+      def set(values)
+        attributes.merge!(values)
+      end
+
+      # increments a number or adds to a set, starts at 0 or [] if it doesn't yet exist
+      def add(values)
+        additions.merge!(values)
+      end
+
+      # deletes a value or values from a set type or simply sets a field to nil
+      def delete(field_or_values)
+        if field_or_values.is_a?(Hash)
+          deletions.merge!(field_or_values)
+        else
+          # adds to array of fields for use in REMOVE update expression
+          removals << field_or_values
+        end
+      end
+
       def find_from_attributes(model_or_model_class, attributes)
         model_class = model_or_model_class.is_a?(Dynamoid::Document) ? model_or_model_class.class : model_or_model_class
         if attributes.is_a?(Hash)

lib/dynamoid/transaction_write/update_upsert.rb

@@ -5,8 +5,8 @@
 module Dynamoid
   class TransactionWrite
     class UpdateUpsert < Action
-      def initialize(model_or_model_class, attributes = {}, options = {})
-        super(model_or_model_class, attributes, options)
+      def initialize(model_or_model_class, attributes = {}, options = {}, &block)
+        super(model_or_model_class, attributes, options, &block)
 
         write_attributes_to_model
       end
@@ -36,20 +36,17 @@ def to_h
 
         # e.g. {":updated_at" => 1645453.234, ":i" => 1}
         expression_attribute_values = item_keys.each_with_index.map { |k, i| [":_s#{i}", item[k]] }.to_h
+        expression_attribute_names = {}
 
-        if options[:add] # not yet documented or supported, may change in a future release
-          # ADD statements can be used to increment a counter:
-          # txn.update!(UserCount, "UserCount#Red", {}, options: {add: {record_count: 1}})
-          add_keys = options[:add].keys
-          update_expression += " ADD #{add_keys.each_with_index.map { |k, i| "#{k} :_a#{i}" }.join(', ')}"
-          add_keys.each_with_index { |k, i| expression_attribute_values[":_a#{i}"] = options[:add][k] }
-        end
+        update_expression = set_additions(expression_attribute_values, update_expression)
+        update_expression = set_deletions(expression_attribute_values, update_expression)
+        expression_attribute_names, update_expression = set_removals(expression_attribute_names, update_expression)
 
         # only alias names for fields in models, other values such as for ADD do not have them
         # e.g. {"#updated_at" => "updated_at"}
         # attribute_keys_in_model = item_keys.intersection(model_class.attributes.keys)
         # expression_attribute_names = attribute_keys_in_model.map{|k| ["##{k}","#{k}"]}.to_h
-        expression_attribute_names = item_keys.each_with_index.map { |k, i| ["#_n#{i}", k.to_s] }.to_h
+        expression_attribute_names.merge!(item_keys.each_with_index.map { |k, i| ["#_n#{i}", k.to_s] }.to_h)
 
         condition_expression = "attribute_exists(#{model_class.hash_key})" # fail if record is missing
         condition_expression += " and attribute_exists(#{model_class.range_key})" if model_class.range_key? # needed?
@@ -67,6 +64,57 @@ def to_h
 
         result
       end
+
+      private
+
+      # adds all of the ADD statements to the update_expression and returns it
+      def set_additions(expression_attribute_values, update_expression)
+        return update_expression unless additions.present?
+
+        # ADD statements can be used to increment a counter:
+        # txn.update!(UserCount, "UserCount#Red", {}, options: {add: {record_count: 1}})
+        add_keys = additions.keys
+        update_expression += " ADD #{add_keys.each_with_index.map { |k, i| "#{k} :_a#{i}" }.join(', ')}"
+        # convert any enumerables into sets
+        add_values = additions.transform_values do |v|
+          if !v.is_a?(Set) && v.is_a?(Enumerable)
+            Set.new(v)
+          else
+            v
+          end
+        end
+        add_keys.each_with_index { |k, i| expression_attribute_values[":_a#{i}"] = add_values[k] }
+        update_expression
+      end
+
+      # adds all of the DELETE statements to the update_expression and returns it
+      def set_deletions(expression_attribute_values, update_expression)
+        return update_expression unless deletions.present?
+
+        delete_keys = deletions.keys
+        update_expression += " DELETE #{delete_keys.each_with_index.map { |k, i| "#{k} :_d#{i}" }.join(', ')}"
+        # values must be sets
+        delete_values = deletions.transform_values do |v|
+          if v.is_a?(Set)
+            v
+          else
+            Set.new(v.is_a?(Enumerable) ? v : [v])
+          end
+        end
+        delete_keys.each_with_index { |k, i| expression_attribute_values[":_d#{i}"] = delete_values[k] }
+        update_expression
+      end
+
+      # adds all of the removals as a REMOVE clause
+      def set_removals(expression_attribute_names, update_expression)
+        return expression_attribute_names, update_expression unless removals.present?
+
+        update_expression += " REMOVE #{removals.each_with_index.map { |_k, i| "#_r#{i}" }.join(', ')}"
+        expression_attribute_names = expression_attribute_names.merge(
+          removals.each_with_index.map { |k, i| ["#_r#{i}", k.to_s] }.to_h
+        )
+        [expression_attribute_names, update_expression]
+      end
     end
   end
 end

spec/dynamoid/transaction_write/context.rb

@@ -8,6 +8,9 @@
   let(:klass) do
     new_class(class_name: 'Document') do
       field :name
+      field :record_count, :integer
+      field :favorite_numbers, :set, of: :integer
+      field :favorite_names, :set, of: :string
     end
   end