local.rb in dynamoid-0.4.0

- old
+ new

@@ -36,11 +36,11 @@
         Hash.new { |h, k| h[k] = Array.new }.tap do |hash|
           options.each do |table_name, keys|
             table = data[table_name]
             if table[:range_key]
               Array(keys).each do |hash_key, range_key|
-                hash[table_name] << get_item(table_name, hash_key, range_key)
+                hash[table_name] << get_item(table_name, hash_key, :range_key => range_key)
               end
             else
               Array(keys).each do |key|
                 hash[table_name] << get_item(table_name, key)
               end
@@ -55,21 +55,23 @@
       # @param [Symbol] key the table's primary key (defaults to :id)
       # @param [Hash] options provide a range_key here if you want one for the table
       #
       # @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.
       #
       # @param [String] table_name the name of the table
       # @param [String] key the hash key of the item to delete
       # @param [Number] range_key the range key of the item to delete, required if the table has a composite key
       #
       # @since 0.2.0
-      def delete_item(table_name, key, range_key = nil)
+      def delete_item(table_name, key, options = {})
+        range_key = options.delete(:range_key)
         data[table_name][:data].delete("#{key}.#{range_key}")
       end
     
       # Deletes an entire table from the hash.
       #
@@ -89,11 +91,12 @@
       # @param [Number] range_key the range key of the item to find, required if the table has a composite key
       #
       # @return [Hash] a hash representing the raw item
       #
       # @since 0.2.0
-      def get_item(table_name, key, range_key = nil)
+      def get_item(table_name, key, options = {})
+        range_key = options[:range_key]
         if data[table_name][:data]
           data[table_name][:data]["#{key}.#{range_key}"]
         else
           nil
         end
@@ -114,10 +117,12 @@
       # @since 0.2.0
       def put_item(table_name, object)
         table = data[table_name]
         table[:data][object[table[:hash_key]]]
         table[:data]["#{object[table[:hash_key]]}.#{object[table[:range_key]]}"] = object.delete_if{|k, v| v.nil? || (v.respond_to?(:empty?) && v.empty?)}
+      rescue
+        raise data.inspect
       end
     
       # Query the hash.
       #
       # @param [String] table_name the name of the table
@@ -132,40 +137,60 @@
       # @return [Array] an array of all matching items
       #
       # @since 0.2.0
       def query(table_name, opts = {})
         id = opts[:hash_value]
+        hash_key = data[table_name][:hash_key]
         range_key = data[table_name][:range_key]
-        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])}
+
+        results = if opts[:range_value]
+          data[table_name][:data].values.find_all{|v| v[hash_key] == 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]}
+          data[table_name][:data].values.find_all{|v| v[hash_key] == id && !v[range_key].nil? && v[range_key] > opts[:range_greater_than]}
         elsif opts[:range_less_than]
-          data[table_name][:data].values.find_all{|v| v[:id] == id && !v[range_key].nil? && v[range_key] < opts[:range_less_than]}
+          data[table_name][:data].values.find_all{|v| v[hash_key] == id && !v[range_key].nil? && v[range_key] < opts[:range_less_than]}
         elsif opts[:range_gte]
-          data[table_name][:data].values.find_all{|v| v[:id] == id && !v[range_key].nil? && v[range_key] >= opts[:range_gte]}
+          data[table_name][:data].values.find_all{|v| v[hash_key] == id && !v[range_key].nil? && v[range_key] >= opts[:range_gte]}
         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]}
+          data[table_name][:data].values.find_all{|v| v[hash_key] == 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[hash_key] == id}
         end
+
+        results = drop_till_start(results, opts[:next_token], range_key, hash_key)
+        results = results.take(opts[:limit]) if opts[:limit]
+        results
       end
     
       # Scan the hash.
       #
       # @param [String] table_name the name of the table
       # @param [Hash] scan_hash a hash of attributes: matching records will be returned by the scan
       #
       # @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], data[table_name][:hash_key])
+        results = results.take(opts[:limit]) if opts[:limit]
+        results
       end
+
+      def drop_till_start(results, next_token, range_key, hash_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[hash_key] != 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
\ No newline at end of file
+end