lib/cassanity/executors/cql_rb.rb in cassanity-0.6.0.beta4 vs lib/cassanity/executors/cql_rb.rb in cassanity-0.6.0.beta5

@@ -146,10 +146,11 @@
             raise Cassanity::UnknownCommand
           end
 
           arguments = args[:arguments]
 
+          send_use_command = false
           if arguments
             # TODO: As a temporary measure, we remove this deprecated option
             # while we have time to update each gem (e.g., adapter-cassanity)
             # that sets it. Consistency should be specified at the connection
             # level for now.
@@ -163,22 +164,25 @@
             end
             if (column_family_name = arguments[:column_family_name])
               payload[:column_family_name] = column_family_name
             end
 
-            # Select keyspace before query runs
+            # Select the correct keyspace before executing the CQL query
             if command != :keyspace_create && (keyspace_name = arguments[:keyspace_name])
-              @driver.use(keyspace_name)
+              send_use_command = true
             end
           end
 
           begin
             cql, *variables = generator.call(arguments)
             payload[:cql] = cql
             payload[:cql_variables] = variables
 
             statement = Cassanity::Statement.new(cql)
-            result = @retry_strategy.execute(payload) { @driver.execute(statement.interpolate(variables)) }
+            result = @retry_strategy.execute(payload) do
+              @driver.use(keyspace_name) if send_use_command
+              @driver.execute(statement.interpolate(variables))
+            end
 
             transformer = @result_transformers.fetch(command, Mirror)
             transformed_result = transformer.call(result, args[:transformer_arguments])
             payload[:result] = transformed_result
           rescue StandardError => e