lib/ddbcli/ddb-driver.rb in ddbcli-0.1.6 vs lib/ddbcli/ddb-driver.rb in ddbcli-0.1.7

@@ -191,11 +191,14 @@
       table_names = do_show_tables0
       h = {}
 
       table_names.map do |table_name|
         table_info = @client.query('DescribeTable', 'TableName' => table_name)['Table']
-        h[table_name] = table_info['TableStatus']
+        h[table_name] = {
+          'TableStatus' => table_info['TableStatus'],
+          'ItemCount'   => table_info['ItemCount'],
+        }
       end
 
       return h
     end
 
@@ -477,21 +480,31 @@
           res_data = select_proc.call(res_data['LastEvaluatedKey'])
           retval += res_data['Count']
         end
       else
         retval = block_given? ? res_data['Items'].map {|i| yield(i) } : res_data['Items']
+        limit_orig = parsed.limit
 
-        if (@iteratable or opts[:iteratable]) and not parsed.limit
+        if @iteratable or opts[:iteratable] or (parsed.limit and retval.length < parsed.limit)
+          parsed.limit -= retval.length if parsed.limit
+
           while res_data['LastEvaluatedKey']
             res_data = select_proc.call(res_data['LastEvaluatedKey'])
-            retval.concat(
-              block_given? ? res_data['Items'].map {|i| yield(i) } : res_data['Items']
-            )
+            items = block_given? ? res_data['Items'].map {|i| yield(i) } : res_data['Items']
+
+            retval.concat(items)
+
+            if parsed.limit
+              parsed.limit -= items.length
+              break if parsed.limit < 1
+            end
           end
         end
       end
 
+      parsed.limit = limit_orig;
+
       if res_data['LastEvaluatedKey']
         @last_action = action
         @last_parsed = parsed
         @last_evaluated_key = res_data['LastEvaluatedKey']
         retval = DynamoDB::Iteratorable.new(retval, res_data['LastEvaluatedKey'])
@@ -665,19 +678,24 @@
             end
           end
         end # scan filter
 
         res_data = @client.query('Scan', req_hash)
-        items.concat(res_data['Items'])
+        res_data_items = res_data['Items']
+        parsed.limit -= res_data_items.length if parsed.limit
+        items.concat(res_data_items)
         res_data['LastEvaluatedKey']
       end
 
       lek = nil
 
       loop do
         lek = scan.call(lek)
-        break unless lek
+
+        if not lek or (parsed.limit and parsed.limit < 1)
+          break
+        end
       end
 
       return items
     end
 
@@ -759,10 +777,16 @@
 
       Rownum.new(n)
     end
 
     def do_insert_select(action, parsed)
+      if parsed.count
+        raise DynamoDB::Error, '"COUNT(*)" cannot be inserted.'
+      end
+
       items = do_select0(action, parsed.select, :iteratable => true)
+      items = items.data if items.kind_of?(Iteratorable)
+
       n = 0
 
       until (chunk = items.slice!(0, MAX_NUMBER_BATCH_PROCESS_ITEMS)).empty?
         operations = []