batch.rb in cassanity-0.6.0

- old
+ new

@@ -2,10 +2,13 @@
 
 module Cassanity
   module ArgumentGenerators
     class Batch
 
+      # Private: List of supported batch types
+      BatchTypes = ['COUNTER','LOGGED','UNLOGGED']
+
       # Private: Map of command to argument generator
       Commands = {
         insert: ColumnFamilyInsert.new,
         update: ColumnFamilyUpdate.new,
         delete: ColumnFamilyDelete.new,
@@ -17,14 +20,18 @@
         @commands = args.fetch(:commands) { Commands }
       end
 
       # Internal
       def call(args = {})
+        type = args[:type].to_s.upcase
+        type = 'LOGGED' if type.empty?
+        raise ArgumentError.new("invalid batch type") unless BatchTypes.include?(type)
+
         using = args[:using]
         modifications_argument = args.fetch(:modifications) { [] }
 
         variables = []
-        cql = "BEGIN BATCH"
+        cql = type == 'LOGGED' ? "BEGIN BATCH" : "BEGIN #{type} BATCH"
 
         using_cql, *using_variables = @using_clause.call(using: using)
         cql << using_cql
         variables.concat(using_variables)
 