Merge pull request #476 from bmalinconico/allow_index_selection

Adding #with_index to chain to force allow the user to force an index

Author: andrykonchin

lib/dynamoid/criteria/chain.rb

@@ -110,7 +110,7 @@ def where(args)
         query.update(args.symbolize_keys)
 
         # we should re-initialize keys detector every time we change query
-        @key_fields_detector = KeyFieldsDetector.new(@query, @source)
+        @key_fields_detector = KeyFieldsDetector.new(@query, @source, forced_index_name: @forced_index_name)
 
         self
       end
@@ -358,6 +358,36 @@ def scan_index_forward(scan_index_forward)
         self
       end
 
+      # Force the index name to use for queries.
+      #
+      # By default allows the library to select the most appropriate index.
+      # Sometimes you have more than one index which will fulfill your query's
+      # needs. When this case occurs you may want to force an order. This occurs
+      # when you are searching by hash key, but not specifying a range key.
+      #
+      #  class Comment
+      #    include Dynamoid::Document
+      #
+      #    table key: :post_id
+      #    range_key :author_id
+      #
+      #    field :post_date, :datetime
+      #
+      #    global_secondary_index name: :time_sorted_comments, hash_key: :post_id, range_key: post_date, projected_attributes: :all
+      #  end
+      #
+      #
+      #   Comment.where(post_id: id).with_index(:time_sorted_comments).scan_index_forward(false)
+      #
+      # @return [Dynamoid::Criteria::Chain]
+      def with_index(index_name)
+        raise Dynamoid::Errors::InvalidIndex, "Unknown index #{index_name}" unless @source.find_index_by_name(index_name)
+
+        @forced_index_name = index_name
+        @key_fields_detector = KeyFieldsDetector.new(@query, @source, forced_index_name: index_name)
+        self
+      end
+
       # Allows to use the results of a search as an enumerable over the results
       # found.
       #

lib/dynamoid/criteria/key_fields_detector.rb

@@ -24,10 +24,11 @@ def contain_with_eq_operator?(field_name)
         end
       end
 
-      def initialize(query, source)
+      def initialize(query, source, forced_index_name: nil)
         @query = query
         @source = source
         @query = Query.new(query)
+        @forced_index_name = forced_index_name
         @result = find_keys_in_query
       end
 
@@ -54,6 +55,8 @@ def index_name
       private
 
       def find_keys_in_query
+        return match_forced_index if @forced_index_name
+
         match_table_and_sort_key ||
           match_local_secondary_index ||
           match_global_secondary_index_and_sort_key ||
@@ -133,6 +136,16 @@ def match_global_secondary_index
           }
         end
       end
+
+      def match_forced_index
+        idx = @source.find_index_by_name(@forced_index_name)
+
+        {
+          hash_key: idx.hash_key,
+          range_key: idx.range_key,
+          index_name: idx.name,
+        }
+      end
     end
   end
 end

lib/dynamoid/indexes.rb

@@ -154,6 +154,16 @@ def find_index(hash, range = nil)
         index
       end
 
+      # Returns an index by its name
+      #
+      # @param name [string, symbol] the name of the index to lookup
+      # @return [Dynamoid::Indexes::Index, nil] index object or nil if it isn't found
+      def find_index_by_name(name)
+        string_name = name.to_s
+        indexes.each_value.detect{ |i| i.name.to_s == string_name }
+      end
+
+
       # Returns true iff the provided hash[,range] key combo is a local
       # secondary index.
       #

spec/dynamoid/criteria/chain_spec.rb

@@ -1803,6 +1803,82 @@ def request_params
     end
   end
 
+  describe '#with_index' do
+    context 'when Local Secondary Index (LSI) used' do
+      let(:klass_with_local_secondary_index) do
+        new_class do
+          range :owner_id
+
+          field :age, :integer
+
+          local_secondary_index range_key: :age,
+            name: :age_index, projected_attributes: :all
+        end
+      end
+
+      before do
+        klass_with_local_secondary_index.create(id: 'the same id', owner_id: 'a', age: 3)
+        klass_with_local_secondary_index.create(id: 'the same id', owner_id: 'c', age: 2)
+        klass_with_local_secondary_index.create(id: 'the same id', owner_id: 'b', age: 1)
+      end
+
+      it 'sorts the results in ascending order' do
+        chain = Dynamoid::Criteria::Chain.new(klass_with_local_secondary_index)
+        models = chain.where(id: 'the same id').with_index(:age_index).scan_index_forward(true)
+        expect(models.map(&:owner_id)).to eq %w[b c a]
+      end
+
+      it 'sorts the results in desc order' do
+        chain = Dynamoid::Criteria::Chain.new(klass_with_local_secondary_index)
+        models = chain.where(id: 'the same id').with_index(:age_index).scan_index_forward(false)
+        expect(models.map(&:owner_id)).to eq %w[a c b]
+      end
+    end
+
+    context 'when Global Secondary Index (GSI) used' do
+      let(:klass_with_global_secondary_index) do
+        new_class do
+          range :owner_id
+
+          field :age, :integer
+
+          global_secondary_index hash_key: :owner_id, range_key: :age,
+            name: :age_index, projected_attributes: :all
+        end
+      end
+
+      before do
+        klass_with_global_secondary_index.create(id: 'the same id', owner_id: 'a', age: 3)
+        klass_with_global_secondary_index.create(id: 'the same id', owner_id: 'c', age: 2)
+        klass_with_global_secondary_index.create(id: 'other id',    owner_id: 'a', age: 1)
+      end
+
+      let(:chain) { Dynamoid::Criteria::Chain.new(klass_with_global_secondary_index)  }
+
+      it 'sorts the results in ascending order' do
+        models = chain.where(owner_id: 'a').with_index(:age_index).scan_index_forward(true)
+        expect(models.map(&:age)).to eq [1, 3]
+      end
+
+      it 'sorts the results in desc order' do
+        models = chain.where(owner_id: 'a').with_index(:age_index).scan_index_forward(false)
+        expect(models.map(&:age)).to eq [3, 1]
+      end
+
+      it 'works with string names' do
+        models = chain.where(owner_id: 'a').with_index('age_index').scan_index_forward(false)
+        expect(models.map(&:age)).to eq [3, 1]
+      end
+
+      it 'raises an error when an unknown index is passed' do
+        expect do
+          chain.where(owner_id: 'a').with_index(:missing_index)
+        end.to raise_error Dynamoid::Errors::InvalidIndex, /Unknown index/
+      end
+    end
+
+  end
+
   describe '#scan_index_forward' do
     let(:klass_with_range_key) do
       new_class do