Merge pull request #29 from ananthakumaran/range_feature

range, limit, start

lib/dynamoid/adapter.rb

@@ -110,12 +110,12 @@ def delete(table, id, range_key = nil)
     # @param [Hash] scan_hash a hash of attributes: matching records will be returned by the scan
     #
     # @since 0.2.0    
-    def scan(table, query)
+    def scan(table, query, opts = {})
       if Dynamoid::Config.partitioning?
-        results = benchmark('Scan', table, query) {adapter.scan(table, query)}
+        results = benchmark('Scan', table, query) {adapter.scan(table, query, opts)}
         result_for_partition(results)
       else
-        adapter.scan(table, query)
+        adapter.scan(table, query, opts)
       end
     end
 

lib/dynamoid/adapter/aws_sdk.rb

@@ -64,7 +64,6 @@ def batch_get_item(options)
       # @since 0.2.0
       def create_table(table_name, key = :id, options = {})
         options[:hash_key] ||= {key.to_sym => :string}
-        options[:range_key] = {options[:range_key].to_sym => :number} if options[:range_key]
         read_capacity = options[:read_capacity] || Dynamoid::Config.read_capacity
         write_capacity = options[:write_capacity] || Dynamoid::Config.write_capacity
         table = @@connection.tables.create(table_name, read_capacity, write_capacity, options)
@@ -180,10 +179,10 @@ def query(table_name, opts = {})
       # @return [Array] an array of all matching items
       #
       # @since 0.2.0
-      def scan(table_name, scan_hash)
+      def scan(table_name, scan_hash, select_opts)
         table = get_table(table_name)
         results = []
-        table.items.where(scan_hash).select do |data|
+        table.items.where(scan_hash).select(select_opts) do |data|
           results << data.attributes.symbolize_keys!
         end
         results

lib/dynamoid/adapter/local.rb

@@ -57,7 +57,8 @@ def batch_get_item(options)
       #
       # @since 0.2.0
       def create_table(table_name, key, options = {})
-        data[table_name] = {:hash_key => key, :range_key => options[:range_key], :data => {}}
+        range_key = options[:range_key] && options[:range_key].keys.first
+        data[table_name] = {:hash_key => key, :range_key => range_key, :data => {}}
       end
 
       # Removes an item from the hash.
@@ -135,7 +136,8 @@ def put_item(table_name, object)
       def query(table_name, opts = {})
         id = opts[:hash_value]
         range_key = data[table_name][:range_key]
-        if opts[:range_value]
+
+        results = if opts[:range_value]
           data[table_name][:data].values.find_all{|v| v[:id] == id && !v[range_key].nil? && opts[:range_value].include?(v[range_key])}
         elsif opts[:range_greater_than]
           data[table_name][:data].values.find_all{|v| v[:id] == id && !v[range_key].nil? && v[range_key] > opts[:range_greater_than]}
@@ -146,8 +148,12 @@ def query(table_name, opts = {})
         elsif opts[:range_lte]  
           data[table_name][:data].values.find_all{|v| v[:id] == id && !v[range_key].nil? && v[range_key] <= opts[:range_lte]}
         else
-          get_item(table_name, id)
+          data[table_name][:data].values.find_all{|v| v[:id] == id}
         end
+
+        results = drop_till_start(results, opts[:next_token], range_key)
+        results = results.take(opts[:limit]) if opts[:limit]
+        results
       end
 
       # Scan the hash.
@@ -158,14 +164,28 @@ def query(table_name, opts = {})
       # @return [Array] an array of all matching items
       #
       # @since 0.2.0
-      def scan(table_name, scan_hash)
+      def scan(table_name, scan_hash, opts = {})
         return [] if data[table_name].nil?
-        data[table_name][:data].values.flatten.select{|d| scan_hash.all?{|k, v| !d[k].nil? && d[k] == v}}
+        results = data[table_name][:data].values.flatten.select{|d| scan_hash.all?{|k, v| !d[k].nil? && d[k] == v}}
+        results = drop_till_start(results, opts[:next_token], data[table_name][:range_key])
+        results = results.take(opts[:limit]) if opts[:limit]
+        results
+      end
+
+      def drop_till_start(results, next_token, range_key)
+        return results unless next_token
+
+        hash_value = next_token[:hash_key_element].values.first
+        range_value = next_token[:range_key_element].values.first if next_token[:range_key_element]
+
+        results = results.drop_while do |r|
+          (r[:id] != hash_value or r[range_key] != range_value)
+        end.drop(1)
       end
 
       # @todo Add an UpdateItem method.
 
       # @todo Add an UpdateTable method.
     end
   end
-end
+end

lib/dynamoid/criteria.rb

@@ -9,7 +9,7 @@ module Criteria
 
     module ClassMethods
 
-      [:where, :all, :first, :each].each do |meth|
+      [:where, :all, :first, :each, :limit, :start].each do |meth|
         # Return a criteria chain in response to a method that will begin or end a chain. For more information, 
         # see Dynamoid::Criteria::Chain.
         #
@@ -26,4 +26,4 @@ module ClassMethods
     end
   end
 
-end
+end

lib/dynamoid/criteria/chain.rb

@@ -6,20 +6,20 @@ module Criteria
     # chain to relation). It is a chainable object that builds up a query and eventually executes it either on an index
     # or by a full table scan.
     class Chain
-      attr_accessor :query, :source, :index, :values
+      attr_accessor :query, :source, :index, :values, :limit, :start
       include Enumerable
-      
+
       # Create a new criteria chain.
       #
       # @param [Class] source the class upon which the ultimate query will be performed.
       def initialize(source)
         @query = {}
         @source = source
       end
-      
-      # The workhorse method of the criteria chain. Each key in the passed in hash will become another criteria that the 
-      # ultimate query must match. A key can either be a symbol or a string, and should be an attribute name or 
-      # an attribute name with a range operator. 
+
+      # The workhorse method of the criteria chain. Each key in the passed in hash will become another criteria that the
+      # ultimate query must match. A key can either be a symbol or a string, and should be an attribute name or
+      # an attribute name with a range operator.
       #
       # @example A simple criteria
       #   where(:name => 'Josh')
@@ -32,7 +32,7 @@ def where(args)
         args.each {|k, v| query[k] = v}
         self
       end
-      
+
       # Returns all the records matching the criteria.
       #
       # @since 0.2.0
@@ -42,35 +42,50 @@ def all
 
       # Returns the first record matching the criteria.
       #
-      # @since 0.2.0      
+      # @since 0.2.0
       def first
-        records.first
+        limit(1).first
+      end
+
+      def limit(limit)
+        @limit = limit
+        records
+      end
+
+      def start(start)
+        @start = start
+        self
       end
 
       # Allows you to use the results of a search as an enumerable over the results found.
       #
-      # @since 0.2.0            
+      # @since 0.2.0
       def each(&block)
         records.each(&block)
       end
-      
+
       private
-      
+
       # The actual records referenced by the association.
       #
       # @return [Array] an array of the found records.
       #
       # @since 0.2.0
       def records
-        return records_with_index if index
-        records_without_index
+        if range?
+          records_with_range
+        elsif index
+          records_with_index
+        else
+          records_without_index
+        end
       end
 
       # If the query matches an index on the associated class, then this method will retrieve results from the index table.
       #
       # @return [Array] an array of the found records.
       #
-      # @since 0.2.0      
+      # @since 0.2.0
       def records_with_index
         ids = if index.range_key?
           Dynamoid::Adapter.query(index.table_name, index_query).collect{|r| r[:ids]}.inject(Set.new) {|set, result| set + result}
@@ -85,67 +100,117 @@ def records_with_index
         if ids.nil? || ids.empty?
           []
         else
-          Array(source.find(ids.to_a))
+          ids = ids.to_a
+
+          if @start
+            ids = ids.drop_while { |id| id != @start.id }.drop(1)
+          end
+
+          ids = ids.take(@limit) if @limit
+          Array(source.find(ids))
         end
       end
-
+
+      def records_with_range
+        Dynamoid::Adapter.query(source.table_name, range_query).collect {|hash| source.new(hash).tap { |r| r.new_record = false } }
+      end
+
       # If the query does not match an index, we'll manually scan the associated table to manually find results.
       #
       # @return [Array] an array of the found records.
       #
-      # @since 0.2.0      
+      # @since 0.2.0
       def records_without_index
         if Dynamoid::Config.warn_on_scan
           Dynamoid.logger.warn 'Queries without an index are forced to use scan and are generally much slower than indexed queries!'
           Dynamoid.logger.warn "You can index this query by adding this to #{source.to_s.downcase}.rb: index [#{source.attributes.sort.collect{|attr| ":#{attr}"}.join(', ')}]"
         end
-        Dynamoid::Adapter.scan(source.table_name, query).collect {|hash| source.new(hash).tap { |r| r.new_record = false } }
+
+        Dynamoid::Adapter.scan(source.table_name, query, query_opts).collect {|hash| source.new(hash).tap { |r| r.new_record = false } }
       end
-      
-      # Format the provided query so that it can be used to query results from DynamoDB. 
+
+      # Format the provided query so that it can be used to query results from DynamoDB.
       #
       # @return [Hash] a hash with keys of :hash_value and :range_value
       #
-      # @since 0.2.0      
+      # @since 0.2.0
       def index_query
         values = index.values(query)
         {}.tap do |hash|
           hash[:hash_value] = values[:hash_value]
           if index.range_key?
             key = query.keys.find{|k| k.to_s.include?('.')}
             if key
-              if query[key].is_a?(Range)
-                hash[:range_value] = query[key]
-              else
-                val = query[key].to_f
-                case key.split('.').last
-                when 'gt'
-                  hash[:range_greater_than] = val
-                when 'lt'
-                  hash[:range_less_than] = val
-                when 'gte'
-                  hash[:range_gte] = val
-                when 'lte'
-                  hash[:range_lte] = val
-                end
-              end
+              hash.merge!(range_hash(key))
             else
               raise Dynamoid::Errors::MissingRangeKey, 'This index requires a range key'
             end
           end
         end
       end
 
+      def range_hash(key)
+        val = query[key]
+
+        return { :range_value => query[key] } if query[key].is_a?(Range)
+
+        case key.split('.').last
+        when 'gt'
+          { :range_greater_than => val.to_f }
+        when 'lt'
+          { :range_less_than  => val.to_f }
+        when 'gte'
+          { :range_gte  => val.to_f }
+        when 'lte'
+          { :range_lte => val.to_f }
+        when 'begins_with'
+          { :range_begins_with => val }
+        end
+      end
+
+      def range_query
+        opts = { :hash_value => query[:id] }
+        if key = query.keys.find { |k| k.to_s.include?('.') }
+          opts.merge!(range_key(key))
+        end
+        opts.merge(query_opts)
+      end
+
       # Return an index that fulfills all the attributes the criteria is querying, or nil if none is found.
       #
-      # @since 0.2.0            
+      # @since 0.2.0
       def index
-        index = source.find_index(query.keys.collect{|k| k.to_s.split('.').first})
+        index = source.find_index(query_keys)
         return nil if index.blank?
         index
       end
+
+      def query_keys
+        query.keys.collect{|k| k.to_s.split('.').first}
+      end
+
+      def range?
+        return false unless source.range_key
+        query_keys == ['id'] || (query_keys.to_set == ['id', source.range_key.to_s].to_set)
+      end
+
+      def start_key
+        key = { :hash_key_element => { 'S' => @start.id } }
+        if range_key = @start.class.range_key
+          range_key_type = @start.class.attributes[range_key][:type] == :string ? 'S' : 'N'
+          key.merge!({:range_key_element => { range_key_type => @start.send(range_key) } })
+        end
+        key
+      end
+
+      def query_opts
+        opts = {}
+        opts[:limit] = @limit if @limit
+        opts[:next_token] = start_key if @start
+        opts
+      end
     end
-    
+
   end
-  
+
 end

lib/dynamoid/fields.rb

@@ -9,6 +9,8 @@ module Fields
     # Initialize the attributes we know the class has, in addition to our magic attributes: id, created_at, and updated_at.
     included do
       class_attribute :attributes
+      class_attribute :range_key
+
       self.attributes = {}
 
       field :id
@@ -36,6 +38,11 @@ def field(name, type = :string, options = {})
 
         respond_to?(:define_attribute_method) ? define_attribute_method(name) : define_attribute_methods([name])
       end
+
+      def range(name, type = :string)
+        field(name, type)
+        self.range_key = name
+      end
     end
 
     # You can access the attributes of an object directly on its attributes method, which is by default an empty hash.

lib/dynamoid/indexes.rb

@@ -40,7 +40,7 @@ def find_index(index)
       # @since 0.2.0
       def create_indexes
         self.indexes.each do |name, index|
-          opts = index.range_key? ? {:range_key => :range} : {}
+          opts = index.range_key? ? {:range_key => { :range => :number }} : {}
           self.create_table(index.table_name, :id, opts) unless self.table_exists?(index.table_name)
         end
       end