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

@@ -50,10 +50,12 @@
       end
 
       retval = case command
                when :SHOW_TABLES
                  do_show_tables(parsed)
+               when :SHOW_TABLE_STATUS
+                 do_show_table_status(parsed)
                when :SHOW_REGIONS
                  do_show_regions(parsed)
                when :SHOW_CREATE_TABLE
                  do_show_create_table(parsed)
                when :ALTER_TABLE
@@ -82,10 +84,14 @@
                  do_delete(parsed)
                when :DELETE_ALL
                  do_delete_all(parsed)
                when :INSERT
                  do_insert(parsed)
+               when :INSERT_SELECT
+                 do_insert_select('Query', parsed)
+               when :INSERT_SCAN
+                 do_insert_select('Scan', parsed)
                when :NEXT
                  if @last_action and @last_parsed and @last_evaluated_key
                    do_select(@last_action, @last_parsed, :last_evaluated_key => @last_evaluated_key)
                  else
                    []
@@ -150,14 +156,18 @@
     end
 
     private
 
     def do_show_tables(parsed)
+      do_show_tables0(parsed.limit)
+    end
+
+    def do_show_tables0(limit = nil)
       req_hash = {}
       table_names = []
 
-      req_hash['Limit'] = parsed.limit if parsed.limit
+      req_hash['Limit'] = limit if limit
 
       list = lambda do |last_evaluated_table_name|
         req_hash['ExclusiveStartTableName'] = last_evaluated_table_name if last_evaluated_table_name
         res_data = @client.query('ListTables', req_hash)
         table_names.concat(res_data['TableNames'])
@@ -167,18 +177,30 @@
       letn = nil
 
       loop do
         letn = list.call(letn)
 
-        if parsed.limit or not letn
+        if limit or not letn
           break
         end
       end
 
       return table_names
     end
 
+    def do_show_table_status(parsed)
+      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']
+      end
+
+      return h
+    end
+
     def do_show_regions(parsed)
       DynamoDB::Endpoint.regions
     end
 
     def do_show_create_table(parsed)
@@ -382,10 +404,16 @@
     def do_describe(parsed)
       (@client.query('DescribeTable', 'TableName' => parsed.table) || {}).fetch('Table', {})
     end
 
     def do_select(action, parsed, opts = {})
+      do_select0(action, parsed, opts) do |i|
+        convert_to_ruby_value(i)
+      end
+    end
+
+    def do_select0(action, parsed, opts = {})
       select_proc = lambda do |last_evaluated_key|
         req_hash = {'TableName' => parsed.table}
         req_hash['AttributesToGet'] = parsed.attrs unless parsed.attrs.empty?
         req_hash['Limit'] = parsed.limit if parsed.limit
         req_hash['ExclusiveStartKey'] = last_evaluated_key if last_evaluated_key
@@ -448,16 +476,18 @@
         while res_data['LastEvaluatedKey']
           res_data = select_proc.call(res_data['LastEvaluatedKey'])
           retval += res_data['Count']
         end
       else
-        retval = res_data['Items'].map {|i| convert_to_ruby_value(i) }
+        retval = block_given? ? res_data['Items'].map {|i| yield(i) } : res_data['Items']
 
-        if @iteratable and not parsed.limit
+        if (@iteratable or opts[:iteratable]) and not parsed.limit
           while res_data['LastEvaluatedKey']
             res_data = select_proc.call(res_data['LastEvaluatedKey'])
-            retval.concat(res_data['Items'].map {|i| convert_to_ruby_value(i) })
+            retval.concat(
+              block_given? ? res_data['Items'].map {|i| yield(i) } : res_data['Items']
+            )
           end
         end
       end
 
       if res_data['LastEvaluatedKey']
@@ -719,9 +749,35 @@
           }
 
           parsed.attrs.zip(val_list).each do |name, val|
             h[name] = convert_to_attribute_value(val)
           end
+        end
+
+        batch_write_item(req_hash)
+        n += chunk.length
+      end
+
+      Rownum.new(n)
+    end
+
+    def do_insert_select(action, parsed)
+      items = do_select0(action, parsed.select, :iteratable => true)
+      n = 0
+
+      until (chunk = items.slice!(0, MAX_NUMBER_BATCH_PROCESS_ITEMS)).empty?
+        operations = []
+
+        req_hash = {
+          'RequestItems' => {
+            parsed.table => operations,
+          },
+        }
+
+        chunk.each do |item|
+          operations << {
+            'PutRequest' => {'Item' => item}
+          }
         end
 
         batch_write_item(req_hash)
         n += chunk.length
       end