chain.rb in dynamoid-3.1.0

- old
+ new

@@ -16,12 +16,13 @@
         @source = source
         @consistent_read = false
         @scan_index_forward = true
 
         # Honor STI and :type field if it presents
-        if @source.attributes.key?(:type)
-          @query[:'type.in'] = @source.deep_subclasses.map(&:name) << @source.name
+        type = @source.inheritance_field
+        if @source.attributes.key?(type)
+          @query[:"#{type}.in"] = @source.deep_subclasses.map(&:name) << @source.name
         end
       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
@@ -49,10 +50,18 @@
       # @since 0.2.0
       def all
         records
       end
 
+      def count
+        if key_present?
+          count_via_query
+        else
+          count_via_scan
+        end
+      end
+
       # Returns the last fetched record matched the criteria
       # Enumerable doesn't implement `last`, only `first`
       # So we have to implement it ourselves
       #
       def last
@@ -159,9 +168,17 @@
         Enumerator.new do |yielder|
           Dynamoid.adapter.scan(source.table_name, scan_query, scan_opts).each do |hash|
             yielder.yield source.from_database(hash)
           end
         end
+      end
+
+      def count_via_query
+        Dynamoid.adapter.query_count(source.table_name, range_query)
+      end
+
+      def count_via_scan
+        Dynamoid.adapter.scan_count(source.table_name, scan_query, scan_opts)
       end
 
       def range_hash(key)
         name, operation = key.to_s.split('.')
         val = type_cast_condition_parameter(name, query[key])