TransactionWrite supports update_fields() with a block for set, add and delete operations on fields

Author: ckhsponge

README.md

@@ -24,8 +24,8 @@ DynamoDB is not like other document-based databases you might know, and
 is very different indeed from relational databases. It sacrifices
 anything beyond the simplest relational queries and transactional
 support to provide a fast, cost-efficient, and highly durable storage
-solution. If your database requires complicated relational queries and
-transaction support, then this modest Gem cannot provide them for you,
+solution. If your database requires complicated relational queries
+then this modest Gem cannot provide them for you
 and neither can DynamoDB. In those cases you would do better to look
 elsewhere for your database needs.
 
@@ -1187,6 +1187,19 @@ Dynamoid::TransactionWrite.execute do |txn|
   # sets the name and title for a user
   # The user is found by id (that equals 1)
   txn.update_fields(User, '1', name: 'bob', title: 'mister')
+
+  # sets the name, increments a count and deletes a field
+  txn.update_fields(User, 1) do |u1| # a User instance is provided
+    u1.set(name: 'bob')
+    u1.add(article_count: 1)
+    u1.delete(:title)
+  end
+
+  # adds to a set of integers and deletes from a set of strings
+  txn.update_fields(User, 2) do |u2|
+    u2.add(friend_ids: [1, 2])
+    u2.delete(child_names: ['bebe'])
+  end
 end
 ```
 

lib/dynamoid/transaction_write.rb

@@ -8,6 +8,7 @@
 require 'dynamoid/transaction_write/update_fields'
 require 'dynamoid/transaction_write/update_attributes'
 require 'dynamoid/transaction_write/upsert'
+require 'dynamoid/transaction_write/item_updater'
 
 module Dynamoid
   # The class +TransactionWrite+ provides means to perform multiple modifying
@@ -345,6 +346,16 @@ def upsert(model_class, hash_key, range_key = nil, attributes) # rubocop:disable
     #     t.update_fields(User, '1', 'Tylor', age: 26)
     #   end
     #
+    # Updates can also be performed in a block.
+    #   Dynamoid::TransactionWrite.execute do |t|
+    #     t.update_fields(User, 1) do |u|
+    #       u.set age: 27, last_name: 'Tylor'
+    #       u.add article_count: 1 # increment a counter
+    #       u.delete favorite_colors: 'green' # remove from a set
+    #       u.delete :first_name # clear a field
+    #     end
+    #   end
+    #
     # Raises a +Dynamoid::Errors::UnknownAttribute+ exception if any of the
     # attributes is not declared in the model class.
     #
@@ -353,8 +364,14 @@ def upsert(model_class, hash_key, range_key = nil, attributes) # rubocop:disable
     # @param range_key [Scalar value] range key value (optional)
     # @param attributes [Hash]
     # @return [nil]
-    def update_fields(model_class, hash_key, range_key = nil, attributes) # rubocop:disable Style/OptionalArguments
+    def update_fields(model_class, hash_key, range_key = nil, attributes = nil)
+      if attributes.nil? && range_key.is_a?(Hash)
+        attributes = range_key
+        range_key = nil
+      end
       action = Dynamoid::TransactionWrite::UpdateFields.new(model_class, hash_key, range_key, attributes)
+      # cannot pass in &block due to code climate max-args limit so instead yield it here
+      yield(action) if block_given?
       register_action action
     end
 

lib/dynamoid/transaction_write/item_updater.rb

@@ -0,0 +1,108 @@
+# frozen_string_literal: true
+
+module Dynamoid
+  class TransactionWrite
+    class ItemUpdater
+      def initialize
+        @additions = {}
+        @deletions = {}
+        @removals = []
+      end
+
+      def empty?
+        @additions.empty? && @deletions.empty? && @removals.empty?
+      end
+
+      # adds to array of fields for use in REMOVE update expression
+      def remove(field)
+        @removals << field
+      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
+          remove(field_or_values)
+        end
+      end
+
+      def merge_update_expression(expression_attribute_names, expression_attribute_values, update_expression)
+        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)
+        [expression_attribute_names, update_expression]
+      end
+
+      # 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
+        add_values = transform_add_values(@additions)
+
+        update_expression += build_add_expression(add_keys)
+        add_keys.each_with_index { |k, i| expression_attribute_values[":_a#{i}"] = add_values[k] }
+        update_expression
+      end
+
+      def transform_add_values(additions)
+        additions.transform_values do |v|
+          !v.is_a?(Set) && v.is_a?(Enumerable) ? Set.new(v) : v
+        end
+      end
+
+      def build_add_expression(add_keys)
+        " ADD #{add_keys.each_with_index.map { |k, i| "#{k} :_a#{i}" }.join(', ')}"
+      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?
+
+        update_expression += build_delete_expression(@deletions.keys)
+        delete_values = prepare_delete_values(@deletions)
+
+        @deletions.keys.each_with_index do |key, index|
+          expression_attribute_values[":_d#{index}"] = delete_values[key]
+        end
+
+        update_expression
+      end
+
+      def build_delete_expression(delete_keys)
+        " DELETE #{delete_keys.each_with_index.map { |key, index| "#{key} :_d#{index}" }.join(', ')}"
+      end
+
+      def prepare_delete_values(deletions)
+        deletions.transform_values(&method(:prepare_delete_values_transform))
+      end
+
+      def prepare_delete_values_transform(value)
+        if value.is_a?(Set)
+          value
+        else
+          Set.new(value.is_a?(Enumerable) ? value : [value])
+        end
+      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

lib/dynamoid/transaction_write/update_fields.rb

@@ -12,7 +12,11 @@ def initialize(model_class, hash_key, range_key, attributes)
         @model_class = model_class
         @hash_key = hash_key
         @range_key = range_key
-        @attributes = attributes
+        @attributes = attributes || {}
+
+        @item_updater = ItemUpdater.new
+
+        yield(self) if block_given?
       end
 
       def on_registration
@@ -29,13 +33,33 @@ def aborted?
       end
 
       def skipped?
-        @attributes.empty?
+        @attributes.empty? && @item_updater.empty?
       end
 
       def observable_by_user_result
         nil
       end
 
+      # sets a value in the attributes
+      def set(attributes)
+        @attributes.merge!(attributes)
+      end
+
+      # adds to array of fields for use in REMOVE update expression
+      def remove(field)
+        @item_updater.remove(field)
+      end
+
+      # increments a number or adds to a set, starts at 0 or [] if it doesn't yet exist
+      def add(values)
+        @item_updater.add(values)
+      end
+
+      # deletes a value or values from a set type or simply sets a field to nil
+      def delete(field_or_values)
+        @item_updater.delete(field_or_values)
+      end
+
       def action_request
         # changed attributes to persist
         changes = @attributes.dup
@@ -63,6 +87,10 @@ def action_request
 
         update_expression = "SET #{update_expression_statements.join(', ')}"
 
+        expression_attribute_names, update_expression = @item_updater.merge_update_expression(
+          expression_attribute_names, expression_attribute_values, update_expression
+        )
+
         # require primary key to exist
         condition_expression = "attribute_exists(#{@model_class.hash_key})"
         if @model_class.range_key?