batch insertion support

lib/cassava/client.rb

@@ -41,6 +41,31 @@ def insert(table, data)
       end
     end
 
+    # @param table [Symbol] the table name
+    # @param data [Hash] A hash of column names to data, which will be inserted into the table
+    # @return [BatchInsertionElement] An object to hold the necessary data for generating a batch insert later
+    def generate_batch_insertion_element(table, data)
+      ttl = data.delete(:ttl)
+      optional_timestamp = data.delete(:optional_timestamp)
+      consistency = data.delete(:consistency)
+
+      statement = insert_statement(table, data, ttl, optional_timestamp)
+
+      BatchInsertionElement.new(statement, data)
+    end
+
+    # @param data [Array] An array of BatchInsertionElements, to be added to the session batch and executed
+    # in a single round trip
+    def batch_insert(batch_inserts_collected)
+      batch_object = session.batch do |batched_statement|
+        batch_inserts_collected.each do |batch_element|
+          batched_statement.add(batch_element.statement, batch_element.arguments)
+        end
+      end
+
+      session.execute(batch_object)
+    end
+
     # @param table [Symbol] the table name
     # @param columns [Array<Symbol>] An optional list of column names (as symbols), to only select those columns
     # @return [StatementBuilder] A statement builder representing the partially completed statement.
@@ -87,7 +112,7 @@ def execute(statement, opts = {})
     def insert_statement(table, data, ttl = nil, optional_timestamp = nil)
       column_names = data.keys
       statement_cql = "INSERT INTO #{table} (#{column_names.join(', ')}) VALUES (#{column_names.map { |x| '?' }.join(',')})"
-      
+
       if ttl && optional_timestamp
         statement_cql += " USING TTL #{ttl} AND TIMESTAMP #{optional_timestamp}" 
       elsif ttl
@@ -325,4 +350,17 @@ def log_error(e, statement, opts)
   class NullLogger
     def method_missing(*); end
   end
+
+  class BatchInsertionElement
+    attr_reader :statement, :data
+
+    def initialize(prepared_statement, data_for_arguments)
+      @statement = prepared_statement
+      @data = data_for_arguments
+    end
+
+    def arguments
+      { :arguments => data.values }
+    end
+  end
 end

test/cassava/client_test.rb

@@ -72,6 +72,36 @@ def string_keys(hash)
         saved_timestamp = @client.select_writetime(:test, :d, { :id => 'i' })
         assert_equal timestamp, saved_timestamp
       end
+
+      context 'batched inserts' do
+        setup do
+          collected_inserts = []
+          ttl = 12345
+          timestamp = Time.now.to_i * 1000000 + Time.now.usec
+          item1 = { :id => 'i', :a => 1, :b => 'b', :c => "'\"item(", :d => 1, :ttl => ttl, :optional_timestamp => timestamp }
+          item2 = { :id => 'j', :a => 1, :b => 'c', :c => "'\"item(", :d => 2, :ttl => ttl, :optional_timestamp => timestamp }
+
+          collected_inserts << @client.generate_batch_insertion_element(:test, item1)
+          collected_inserts << @client.generate_batch_insertion_element(:test, item2)
+        end
+
+        should 'connect inserts for batching' do
+          collected_inserts = []
+          ttl = 12345
+          timestamp = Time.now.to_i * 1000000 + Time.now.usec
+          item1 = { :id => 'i', :a => 1, :b => 'b', :c => "'\"item(", :d => 1, :ttl => ttl, :optional_timestamp => timestamp }
+          item2 = { :id => 'j', :a => 1, :b => 'c', :c => "'\"item(", :d => 2, :ttl => ttl, :optional_timestamp => timestamp }
+
+          collected_inserts << @client.generate_batch_insertion_element(:test, item1)
+          collected_inserts << @client.generate_batch_insertion_element(:test, item2)
+
+          @client.batch_insert(collected_inserts)
+
+          result = @client.select(:test).execute.rows
+          assert_equal string_keys(item2), result.next
+          assert_equal string_keys(item1), result.next
+        end
+      end
     end
 
     context 'select' do