lib/cassanity/argument_generators/batch.rb in cassanity-0.3.0 vs lib/cassanity/argument_generators/batch.rb in cassanity-0.4.0

@@ -32,9 +32,17 @@
         modifications = []
         modifications_argument.each do |modification|
           command_name, command_arguments = modification
           command = @commands.fetch(command_name)
 
+          if args[:column_family_name]
+            command_arguments[:column_family_name] ||= args[:column_family_name]
+          end
+
+          if args[:keyspace_name]
+            command_arguments[:keyspace_name] ||= args[:keyspace_name]
+          end
+
           modification_cql, *modification_variables = command.call(command_arguments)
           modifications << modification_cql
           variables.concat(modification_variables)
         end