diff --git a/lib/cassava/client.rb b/lib/cassava/client.rb index 4e2f33b..8aaae01 100644 --- a/lib/cassava/client.rb +++ b/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] 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 diff --git a/test/cassava/client_test.rb b/test/cassava/client_test.rb index af3a17b..89c78aa 100644 --- a/test/cassava/client_test.rb +++ b/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