Add specs and refactor implementation

Author: andrykonchin

lib/dynamoid/transaction_write.rb

@@ -347,15 +347,51 @@ def upsert(model_class, hash_key, range_key = nil, attributes) # rubocop:disable
     #   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
+    #       u.add(article_count: 1)
+    #       u.delete(favorite_colors: 'green')
+    #       u.set(age: 27, last_name: 'Tylor')
     #     end
     #   end
     #
+    # Operation +add+ just adds a value for numeric attributes and join
+    # collections if attribute is a set.
+    #
+    #   t.update_fields(User, 1) do |u|
+    #     u.add(age: 1, followers_count: 5)
+    #     u.add(hobbies: ['skying', 'climbing'])
+    #   end
+    #
+    # Operation +delete+ is applied to collection attribute types and
+    # substructs one collection from another.
+    #
+    #   t.update_fields(User, 1) do |u|
+    #     u.delete(hobbies: ['skying'])
+    #   end
+    #
+    # Operation +set+ just changes an attribute value:
+    #
+    #   t.update_fields(User, 1) do |u|
+    #     u.set(age: 21)
+    #   end
+    #
+    # Operation +remove+ removes one or more attributes from an item.
+    #
+    #   t.update_fields(User, 1) do |u|
+    #     u.remove(:age)
+    #   end
+    #
+    # All the operations work like +ADD+, +DELETE+, +REMOVE+, and +SET+ actions supported
+    # by +UpdateExpression+
+    # {parameter}[https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Expressions.UpdateExpressions.html]
+    # of +UpdateItem+ operation.
+    #
+    # It's atomic operations. So adding or deleting elements in a collection
+    # or incrementing or decrementing a numeric field is atomic and does not
+    # interfere with other write requests.
+    #
     # Raises a +Dynamoid::Errors::UnknownAttribute+ exception if any of the
     # attributes is not declared in the model class.
     #
@@ -364,14 +400,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 = nil)
-      if attributes.nil? && range_key.is_a?(Hash)
+    def update_fields(model_class, hash_key, range_key = nil, attributes = nil, &block)
+      # given no attributes, but there may be a block
+      if range_key.is_a?(Hash) && !attributes
         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?
+
+      action = Dynamoid::TransactionWrite::UpdateFields.new(model_class, hash_key, range_key, attributes, &block)
       register_action action
     end
 

lib/dynamoid/transaction_write/item_updater.rb

@@ -3,106 +3,53 @@
 module Dynamoid
   class TransactionWrite
     class ItemUpdater
-      def initialize
-        @additions = {}
-        @deletions = {}
-        @removals = []
-      end
+      attr_reader :attributes_to_set, :attributes_to_add, :attributes_to_delete, :attributes_to_remove
 
-      def empty?
-        @additions.empty? && @deletions.empty? && @removals.empty?
-      end
+      def initialize(model_class)
+        @model_class = model_class
 
-      # 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)
+        @attributes_to_set = {}
+        @attributes_to_add = {}
+        @attributes_to_delete = {}
+        @attributes_to_remove = []
       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
+      def empty?
+        [@attributes_to_set, @attributes_to_add, @attributes_to_delete, @attributes_to_remove].all?(&:empty?)
       end
 
-      def transform_add_values(additions)
-        additions.transform_values do |v|
-          !v.is_a?(Set) && v.is_a?(Enumerable) ? Set.new(v) : v
-        end
+      def set(attributes)
+        validate_attribute_names!(attributes.keys)
+        @attributes_to_set.merge!(attributes)
       end
 
-      def build_add_expression(add_keys)
-        " ADD #{add_keys.each_with_index.map { |k, i| "#{k} :_a#{i}" }.join(', ')}"
+      # adds to array of fields for use in REMOVE update expression
+      def remove(*names)
+        validate_attribute_names!(names)
+        @attributes_to_remove += names
       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
+      # increments a number or adds to a set, starts at 0 or [] if it doesn't yet exist
+      def add(attributes)
+        validate_attribute_names!(attributes.keys)
+        @attributes_to_add.merge!(attributes)
       end
 
-      def build_delete_expression(delete_keys)
-        " DELETE #{delete_keys.each_with_index.map { |key, index| "#{key} :_d#{index}" }.join(', ')}"
+      # deletes a value or values from a set
+      def delete(attributes)
+        validate_attribute_names!(attributes.keys)
+        @attributes_to_delete.merge!(attributes)
       end
 
-      def prepare_delete_values(deletions)
-        deletions.transform_values(&method(:prepare_delete_values_transform))
-      end
+      private
 
-      def prepare_delete_values_transform(value)
-        if value.is_a?(Set)
-          value
-        else
-          Set.new(value.is_a?(Enumerable) ? value : [value])
+      def validate_attribute_names!(names)
+        names.each do |name|
+          unless @model_class.attributes[name]
+            raise Dynamoid::Errors::UnknownAttribute.new(@model_class, name)
+          end
         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