Merge pull request #913 from Dynamoid/ak/fix-sending-partition-key-and-always-dump-it

Fix partition keys with Dynamoid-specific type and always dump them into a proper DynamoDB type

Author: andrykonchin

lib/dynamoid/finders.rb

@@ -109,12 +109,16 @@ def find_by_id(id, options = {})
       def _find_all(ids, options = {})
         raise Errors::MissingRangeKey if range_key && ids.any? { |_pk, sk| sk.nil? }
 
-        if range_key
-          ids = ids.map do |pk, sk|
-            sk_casted = TypeCasting.cast_field(sk, attributes[range_key])
-            sk_dumped = Dumping.dump_field(sk_casted, attributes[range_key])
-
-            [pk, sk_dumped]
+        ids = ids.map do |id|
+          if range_key
+            # expect [hash key, range key] pair
+            pk, sk = id
+            pk_dumped = cast_and_dump(hash_key, pk)
+            sk_dumped = cast_and_dump(range_key, sk)
+
+            [pk_dumped, sk_dumped]
+          else
+            cast_and_dump(hash_key, id)
           end
         end
 
@@ -155,15 +159,13 @@ def _find_all(ids, options = {})
       def _find_by_id(id, options = {})
         raise Errors::MissingRangeKey if range_key && options[:range_key].nil?
 
-        if range_key
-          key = options[:range_key]
-          key_casted = TypeCasting.cast_field(key, attributes[range_key])
-          key_dumped = Dumping.dump_field(key_casted, attributes[range_key])
+        partition_key_dumped = cast_and_dump(hash_key, id)
 
-          options[:range_key] = key_dumped
+        if range_key
+          options[:range_key] = cast_and_dump(range_key, options[:range_key])
         end
 
-        if item = Dynamoid.adapter.read(table_name, id, options.slice(:range_key, :consistent_read))
+        if item = Dynamoid.adapter.read(table_name, partition_key_dumped, options.slice(:range_key, :consistent_read))
           model = from_database(item)
           model.run_callbacks :find
           model
@@ -308,6 +310,14 @@ def method_missing(method, *args)
           super
         end
       end
+
+      private
+
+      def cast_and_dump(name, value)
+        attribute_options = attributes[name]
+        casted_value = TypeCasting.cast_field(value, attribute_options)
+        Dumping.dump_field(casted_value, attribute_options)
+      end
     end
   end
 end

lib/dynamoid/persistence.rb

@@ -685,9 +685,10 @@ def update!(conditions = {})
 
         begin
           table_name = self.class.table_name
+          partition_key_dumped = Dumping.dump_field(hash_key, self.class.attributes[self.class.hash_key])
           update_item_options = options.merge(conditions: conditions)
 
-          new_attrs = Dynamoid.adapter.update_item(table_name, hash_key, update_item_options) do |t|
+          new_attrs = Dynamoid.adapter.update_item(table_name, partition_key_dumped, update_item_options) do |t|
             item_updater = ItemUpdaterWithDumping.new(self.class, t)
 
             item_updater.add(lock_version: 1) if self.class.attributes[:lock_version]
@@ -926,6 +927,7 @@ def destroy!
     # @since 0.2.0
     def delete
       options = range_key ? { range_key: Dumping.dump_field(read_attribute(range_key), self.class.attributes[range_key]) } : {}
+      partition_key_dumped = Dumping.dump_field(hash_key, self.class.attributes[self.class.hash_key])
 
       # Add an optimistic locking check if the lock_version column exists
       if self.class.attributes[:lock_version]
@@ -941,7 +943,7 @@ def delete
 
       @destroyed = true
 
-      Dynamoid.adapter.delete(self.class.table_name, hash_key, options)
+      Dynamoid.adapter.delete(self.class.table_name, partition_key_dumped, options)
 
       self.class.associations.each_key do |name|
         send(name).disassociate_source

lib/dynamoid/persistence/inc.rb

@@ -21,8 +21,9 @@ def initialize(model_class, hash_key, range_key = nil, counters)
 
       def call
         touch = @counters.delete(:touch)
+        hash_key_dumped = cast_and_dump(@model_class.hash_key, @hash_key)
 
-        Dynamoid.adapter.update_item(@model_class.table_name, @hash_key, update_item_options) do |t|
+        Dynamoid.adapter.update_item(@model_class.table_name, hash_key_dumped, update_item_options) do |t|
           item_updater = ItemUpdaterWithCastingAndDumping.new(@model_class, t)
 
           @counters.each do |name, value|
@@ -43,10 +44,8 @@ def call
 
       def update_item_options
         if @model_class.range_key
-          range_key_options = @model_class.attributes[@model_class.range_key]
-          value_casted = TypeCasting.cast_field(@range_key, range_key_options)
-          value_dumped = Dumping.dump_field(value_casted, range_key_options)
-          { range_key: value_dumped }
+          range_key_dumped = cast_and_dump(@model_class.range_key, @range_key)
+          { range_key: range_key_dumped }
         else
           {}
         end
@@ -60,6 +59,12 @@ def timestamp_attributes_to_touch(touch)
         names += Array.wrap(touch) if touch != true
         names
       end
+
+      def cast_and_dump(name, value)
+        options = @model_class.attributes[name]
+        value_casted = TypeCasting.cast_field(value, options)
+        Dumping.dump_field(value_casted, options)
+      end
     end
   end
 end

lib/dynamoid/persistence/save.rb

@@ -36,8 +36,10 @@ def call
           Dynamoid.adapter.write(@model.class.table_name, attributes_dumped, conditions_for_write)
         else
           attributes_to_persist = @model.attributes.slice(*@model.changed.map(&:to_sym))
+          partition_key_dumped = dump(@model.class.hash_key, @model.hash_key)
+          options = options_to_update_item(partition_key_dumped)
 
-          Dynamoid.adapter.update_item(@model.class.table_name, @model.hash_key, options_to_update_item) do |t|
+          Dynamoid.adapter.update_item(@model.class.table_name, partition_key_dumped, options) do |t|
             item_updater = ItemUpdaterWithDumping.new(@model.class, t)
 
             attributes_to_persist.each do |name, value|
@@ -63,11 +65,9 @@ def conditions_for_write
         conditions = {}
 
         # Add an 'exists' check to prevent overwriting existing records with new ones
-        if @model.new_record?
-          conditions[:unless_exists] = [@model.class.hash_key]
-          if @model.range_key
-            conditions[:unless_exists] << @model.range_key
-          end
+        conditions[:unless_exists] = [@model.class.hash_key]
+        if @model.range_key
+          conditions[:unless_exists] << @model.range_key
         end
 
         # Add an optimistic locking check if the lock_version column exists
@@ -81,17 +81,17 @@ def conditions_for_write
         conditions
       end
 
-      def options_to_update_item
+      def options_to_update_item(partition_key_dumped)
         options = {}
 
         if @model.class.range_key
-          value_dumped = Dumping.dump_field(@model.range_value, @model.class.attributes[@model.class.range_key])
+          value_dumped = dump(@model.class.range_key, @model.range_value)
           options[:range_key] = value_dumped
         end
 
         conditions = {}
         conditions[:if] ||= {}
-        conditions[:if][@model.class.hash_key] = @model.hash_key
+        conditions[:if][@model.class.hash_key] = partition_key_dumped
 
         # Add an optimistic locking check if the lock_version column exists
         # Uses the original lock_version value from Dirty API
@@ -105,6 +105,11 @@ def options_to_update_item
 
         options
       end
+
+      def dump(name, value)
+        options = @model.class.attributes[name]
+        Dumping.dump_field(value, options)
+      end
     end
   end
 end

lib/dynamoid/persistence/update_fields.rb

@@ -16,6 +16,8 @@ def initialize(model_class, partition_key:, sort_key:, attributes:, conditions:)
         @sort_key = sort_key
         @attributes = attributes.symbolize_keys
         @conditions = conditions
+
+        @partition_key_dumped = cast_and_dump(@model_class.hash_key, @partition_key)
       end
 
       def call
@@ -33,7 +35,7 @@ def call
       private
 
       def update_item
-        Dynamoid.adapter.update_item(@model_class.table_name, @partition_key, options_to_update_item) do |t|
+        Dynamoid.adapter.update_item(@model_class.table_name, @partition_key_dumped, options_to_update_item) do |t|
           item_updater = ItemUpdaterWithCastingAndDumping.new(@model_class, t)
 
           @attributes.each do |k, v|
@@ -50,18 +52,23 @@ def options_to_update_item
         options = {}
 
         if @model_class.range_key
-          value_casted = TypeCasting.cast_field(@sort_key, @model_class.attributes[@model_class.range_key])
-          value_dumped = Dumping.dump_field(value_casted, @model_class.attributes[@model_class.range_key])
-          options[:range_key] = value_dumped
+          range_key_dumped = cast_and_dump(@model_class.range_key, @sort_key)
+          options[:range_key] = range_key_dumped
         end
 
         conditions = @conditions.deep_dup
         conditions[:if] ||= {}
-        conditions[:if][@model_class.hash_key] = @partition_key
+        conditions[:if][@model_class.hash_key] = @partition_key_dumped
         options[:conditions] = conditions
 
         options
       end
+
+      def cast_and_dump(name, value)
+        options = @model_class.attributes[name]
+        value_casted = TypeCasting.cast_field(value, options)
+        Dumping.dump_field(value_casted, options)
+      end
     end
   end
 end

lib/dynamoid/persistence/upsert.rb

@@ -33,7 +33,9 @@ def call
       private
 
       def update_item
-        Dynamoid.adapter.update_item(@model_class.table_name, @partition_key, options_to_update_item) do |t|
+        partition_key_dumped = cast_and_dump(@model_class.hash_key, @partition_key)
+
+        Dynamoid.adapter.update_item(@model_class.table_name, partition_key_dumped, options_to_update_item) do |t|
           item_updater = ItemUpdaterWithCastingAndDumping.new(@model_class, t)
 
           @attributes.each do |k, v|
@@ -46,9 +48,8 @@ def options_to_update_item
         options = {}
 
         if @model_class.range_key
-          value_casted = TypeCasting.cast_field(@sort_key, @model_class.attributes[@model_class.range_key])
-          value_dumped = Dumping.dump_field(value_casted, @model_class.attributes[@model_class.range_key])
-          options[:range_key] = value_dumped
+          range_key_dumped = cast_and_dump(@model_class.range_key, @sort_key)
+          options[:range_key] = range_key_dumped
         end
 
         options[:conditions] = @conditions
@@ -58,6 +59,12 @@ def options_to_update_item
       def undump_attributes(raw_attributes)
         Undumping.undump_attributes(raw_attributes, @model_class.attributes)
       end
+
+      def cast_and_dump(name, value)
+        options = @model_class.attributes[name]
+        value_casted = TypeCasting.cast_field(value, options)
+        Dumping.dump_field(value_casted, options)
+      end
     end
   end
 end

spec/dynamoid/criteria/chain_spec.rb

@@ -256,6 +256,36 @@ def request_params
         model.where(name: 'Bob', 'age.foo': 10).to_a
       end.to raise_error(Dynamoid::Errors::Error, 'Unsupported operator foo in age.foo')
     end
+
+    context 'primary key dumping' do
+      it 'uses dumped value of partition key to query item' do
+        klass = new_class(partition_key: { name: :published_on, type: :date })
+
+        obj1 = klass.create(published_on: Date.today + 1)
+        obj2 = klass.create(published_on: Date.today + 2)
+
+        chain = described_class.new(klass)
+        expect(chain).to receive(:raw_pages_via_query).and_call_original
+
+        objects_found = chain.where(published_on: obj1.published_on).all
+        expect(objects_found).to contain_exactly(obj1)
+      end
+
+      it 'uses dumped value of sort key to query item' do
+        klass = new_class do
+          range :published_on, :date
+        end
+
+        obj1 = klass.create(published_on: Date.today + 1)
+        obj2 = klass.create(published_on: Date.today + 2)
+
+        chain = described_class.new(klass)
+        expect(chain).to receive(:raw_pages_via_query).and_call_original
+
+        objects_found = chain.where(id: obj1.id, published_on: obj1.published_on).all
+        expect(objects_found).to contain_exactly(obj1)
+      end
+    end
   end
 
   # http://docs.aws.amazon.com/amazondynamodb/latest/developerguide/LegacyConditionalParameters.QueryFilter.html

spec/dynamoid/finders_spec.rb

@@ -3,8 +3,6 @@
 require 'spec_helper'
 
 describe Dynamoid::Finders do
-  let!(:address) { Address.create(city: 'Chicago') }
-
   describe '.find' do
     let(:klass) do
       new_class(class_name: 'Document')
@@ -111,6 +109,22 @@
           end
         end
       end
+
+      it 'uses dumped value of partition key to update item' do
+        klass = new_class(partition_key: { name: :published_on, type: :date })
+
+        obj = klass.create!(published_on: '2018-10-07'.to_date)
+        expect(klass.find(obj.published_on)).to eql(obj)
+      end
+
+      it 'uses dumped value of sort key to update item' do
+        klass = new_class do
+          range :published_on, :date
+        end
+
+        obj = klass.create!(published_on: '2018-10-07'.to_date)
+        expect(klass.find(obj.id, range_key: obj.published_on)).to eql(obj)
+      end
     end
 
     context 'multiple primary keys provided' do
@@ -270,6 +284,26 @@
         end
       end
 
+      it 'uses dumped value of partition key to update item' do
+        klass = new_class(partition_key: { name: :published_on, type: :date })
+        objects = (1..2).map { |i| klass.create(published_on: '2018-10-07'.to_date + i) }
+        obj1, obj2 = objects
+
+        expect(klass.find([obj1.published_on, obj2.published_on])).to match_array(objects)
+      end
+
+      it 'uses dumped value of sort key to update item' do
+        klass = new_class do
+          range :published_on, :date
+        end
+
+        objects = (1..2).map { |i| klass.create(published_on: '2018-10-07'.to_date + i) }
+        obj1, obj2 = objects
+        found_objects = klass.find([[obj1.id, obj1.published_on], [obj2.id, obj2.published_on]])
+
+        expect(found_objects).to match_array(objects)
+      end
+
       context 'field is not declared in document' do
         let(:class_with_not_declared_field) do
           new_class do
@@ -386,6 +420,8 @@
   end
 
   it 'sends consistent option to the adapter' do
+    address = Address.create(city: 'Chicago')
+
     expect(Dynamoid.adapter).to receive(:get_item)
       .with(anything, anything, hash_including(consistent_read: true))
       .and_call_original
@@ -394,6 +430,7 @@
 
   context 'with users' do
     it 'finds using method_missing for attributes' do
+      address = Address.create(city: 'Chicago')
       array = Address.find_by_city('Chicago')
 
       expect(array).to eq address