cql-rb 1.2.2 → 2.0.0.pre0

data/lib/cql/client/asynchronous_client.rb

@@ -5,25 +5,26 @@ module Cql
     # @private
     class AsynchronousClient < Client
       def initialize(options={})
-        compressor = options[:compressor]
+        @compressor = options[:compressor]
+        @cql_version = options[:cql_version]
         @logger = options[:logger] || NullLogger.new
-        @io_reactor = options[:io_reactor] || Io::IoReactor.new(protocol_handler_factory(compressor))
+        @protocol_version = options[:protocol_version] || 2
+        @io_reactor = options[:io_reactor] || Io::IoReactor.new(protocol_handler_factory)
         @hosts = extract_hosts(options)
         @initial_keyspace = options[:keyspace]
+        @connections_per_node = options[:connections_per_node] || 1
         @lock = Mutex.new
-        @connected = false
-        @connecting = false
-        @closing = false
         @request_runner = RequestRunner.new
         @keyspace_changer = KeyspaceChanger.new
         @connection_manager = ConnectionManager.new
-        port = options[:port] || DEFAULT_PORT
-        credentials = options[:credentials]
-        connections_per_node = options[:connections_per_node] || 1
-        connection_timeout = options[:connection_timeout] || DEFAULT_CONNECTION_TIMEOUT
-        default_consistency = options[:default_consistency] || DEFAULT_CONSISTENCY
-        @execute_options_decoder = ExecuteOptionsDecoder.new(default_consistency)
-        @connection_helper = ConnectionHelper.new(@io_reactor, port, credentials, connections_per_node, connection_timeout, compressor, @logger)
+        @execute_options_decoder = ExecuteOptionsDecoder.new(options[:default_consistency] || DEFAULT_CONSISTENCY)
+        @port = options[:port] || DEFAULT_PORT
+        @connection_timeout = options[:connection_timeout] || DEFAULT_CONNECTION_TIMEOUT
+        @credentials = options[:credentials]
+        @auth_provider = options[:auth_provider] || @credentials && PlainTextAuthProvider.new(*@credentials.values_at(:username, :password))
+        @connected = false
+        @connecting = false
+        @closing = false
       end
 
       def connect
@@ -32,12 +33,15 @@ module Cql
           return @connected_future if can_execute?
           @connecting = true
           @connected_future = begin
-            f = @connection_helper.connect(@hosts, @initial_keyspace)
-            f.on_value do |connections|
+            f = @io_reactor.start
+            f = f.flat_map { connect_with_protocol_version_fallback }
+            f = f.flat_map { |connections| connect_to_all_peers(connections) }
+            f = f.flat_map do |connections|
               @connection_manager.add_connections(connections)
               register_event_listener(@connection_manager.random_connection)
             end
-            f.map { self }
+            f = f.flat_map { use_keyspace(@connection_manager.snapshot, @initial_keyspace) }
+            f.map(self)
           end
         end
         @connected_future.on_complete(&method(:connected))
@@ -49,15 +53,16 @@ module Cql
           return @closed_future if @closing
           @closing = true
           @closed_future = begin
-            f = @io_reactor.stop
             if @connecting
-              ff = @connected_future
-              ff = f.flat_map { ff }
-              ff = f.fallback { ff }
+              f = @connected_future.recover
+              f = f.flat_map { @io_reactor.stop }
+              f = f.map(self)
+              f
             else
-              ff = f
+              f = @io_reactor.stop
+              f = f.map(self)
+              f
             end
-            ff.map { self }
           end
         end
         @closed_future.on_complete(&method(:closed))
@@ -74,20 +79,25 @@ module Cql
 
       def use(keyspace)
         with_failure_handler do
-          connections = @connection_manager.select { |c| c.keyspace != keyspace }
-          if connections.any?
-            futures = connections.map { |connection| @keyspace_changer.use_keyspace(keyspace, connection) }
-            Future.all(*futures).map { nil }
-          else
-            Future.resolved
-          end
+          connections = @connection_manager.reject { |c| c.keyspace == keyspace }
+          return Future.resolved if connections.empty?
+          use_keyspace(connections, keyspace).map(nil)
         end
       end
 
-      def execute(cql, options_or_consistency=nil)
+      def execute(cql, *args)
         with_failure_handler do
-          consistency, timeout, trace = @execute_options_decoder.decode_options(options_or_consistency)
-          execute_request(Protocol::QueryRequest.new(cql, consistency, trace), timeout)
+          options_or_consistency = nil
+          if args.last.is_a?(Symbol) || args.last.is_a?(Hash)
+            options_or_consistency = args.pop
+          end
+          options = @execute_options_decoder.decode_options(options_or_consistency)
+          request = Protocol::QueryRequest.new(cql, args, options[:type_hints], options[:consistency], options[:serial_consistency], options[:page_size], options[:paging_state], options[:trace])
+          f = execute_request(request, options[:timeout])
+          if options.include?(:page_size)
+            f = f.map { |result| AsynchronousQueryPagedQueryResult.new(self, request, result, options) }
+          end
+          f
         end
       end
 
@@ -97,15 +107,32 @@ module Cql
         end
       end
 
+      def batch(type=:logged, options=nil)
+        if type.is_a?(Hash)
+          options = type
+          type = :logged
+        end
+        b = AsynchronousBatch.new(type, @execute_options_decoder, @connection_manager, options)
+        if block_given?
+          yield b
+          b.execute
+        else
+          b
+        end
+      end
+
       private
 
+      DEFAULT_CQL_VERSIONS = {1 => '3.0.0'}
+      DEFAULT_CQL_VERSIONS.default = '3.1.0'
+      DEFAULT_CQL_VERSIONS.freeze
       DEFAULT_CONSISTENCY = :quorum
       DEFAULT_PORT = 9042
       DEFAULT_CONNECTION_TIMEOUT = 10
       MAX_RECONNECTION_ATTEMPTS = 5
 
-      def protocol_handler_factory(compressor)
-        lambda { |connection, timeout| Protocol::CqlProtocolHandler.new(connection, timeout, compressor) }
+      def protocol_handler_factory
+        lambda { |connection, timeout| Protocol::CqlProtocolHandler.new(connection, timeout, @protocol_version, @compressor) }
       end
 
       def extract_hosts(options)
@@ -118,6 +145,52 @@ module Cql
         end
       end
 
+      def create_cluster_connector
+        cql_version = @cql_version || DEFAULT_CQL_VERSIONS[@protocol_version]
+        authentication_step = @protocol_version == 1 ? CredentialsAuthenticationStep.new(@credentials) : SaslAuthenticationStep.new(@auth_provider)
+        ClusterConnector.new(
+          Connector.new([
+            ConnectStep.new(@io_reactor, @port, @connection_timeout, @logger),
+            CacheOptionsStep.new,
+            InitializeStep.new(cql_version, @compressor, @logger),
+            authentication_step,
+            CachePropertiesStep.new,
+          ]),
+          @logger
+        )
+      end
+
+      def connect_with_protocol_version_fallback
+        f = create_cluster_connector.connect_all(@hosts, @connections_per_node)
+        f.fallback do |error|
+          if error.is_a?(QueryError) && error.code == 0x0a && @protocol_version > 1
+            @logger.warn('Could not connect using protocol version %d (will try again with %d): %s' % [@protocol_version, @protocol_version - 1, error.message])
+            @protocol_version -= 1
+            connect_with_protocol_version_fallback
+          else
+            raise error
+          end
+        end
+      end
+
+      def connect_to_all_peers(seed_connections, initial_keyspace=@initial_keyspace)
+        @logger.debug('Looking for additional nodes')
+        peer_discovery = PeerDiscovery.new(seed_connections)
+        peer_discovery.new_hosts.flat_map do |hosts|
+          if hosts.empty?
+            @logger.debug('No additional nodes found')
+            Future.resolved(seed_connections)
+          else
+            @logger.debug('%d additional nodes found' % hosts.size)
+            f = create_cluster_connector.connect_all(hosts, @connections_per_node)
+            f = f.map do |discovered_connections|
+              seed_connections + discovered_connections
+            end
+            f.recover(seed_connections)
+          end
+        end
+      end
+
       def connected(f)
         if f.resolved?
           @lock.synchronize do
@@ -163,40 +236,51 @@ module Cql
         Future.failed(e)
       end
 
+      def use_keyspace(connections, keyspace)
+        futures = connections.map { |connection| @keyspace_changer.use_keyspace(connection, keyspace) }
+        Future.all(*futures)
+      end
+
       def register_event_listener(connection)
         register_request = Protocol::RegisterRequest.new(Protocol::TopologyChangeEventResponse::TYPE, Protocol::StatusChangeEventResponse::TYPE)
-        execute_request(register_request, nil, connection)
-        connection.on_closed do
-          if connected?
-            begin
-              register_event_listener(@connection_manager.random_connection)
-            rescue NotConnectedError
-              # we had started closing down after the connection check
+        f = execute_request(register_request, nil, connection)
+        f.on_value do
+          connection.on_closed do
+            if connected?
+              begin
+                register_event_listener(@connection_manager.random_connection)
+              rescue NotConnectedError
+                # we had started closing down after the connection check
+              end
             end
           end
-        end
-        connection.on_event do |event|
-          if event.change == 'UP' || event.change == 'NEW_NODE'
-            @logger.debug('Received %s event' % event.change)
-            unless @looking_for_nodes
-              @looking_for_nodes = true
-              handle_topology_change.on_complete do |f|
-                @looking_for_nodes = false
+          connection.on_event do |event|
+            if event.change == 'UP' || event.change == 'NEW_NODE'
+              @logger.debug('Received %s event' % event.change)
+              unless @looking_for_nodes
+                @looking_for_nodes = true
+                handle_topology_change.on_complete do |f|
+                  @looking_for_nodes = false
+                end
               end
             end
           end
         end
+        f
       end
 
       def handle_topology_change(remaning_attempts=MAX_RECONNECTION_ATTEMPTS)
         with_failure_handler do
           seed_connections = @connection_manager.snapshot
-          f = @connection_helper.discover_peers(seed_connections, keyspace)
-          f.flat_map do |connections|
-            connected_connections = connections.select(&:connected?)
-            if connected_connections.any?
-              @connection_manager.add_connections(connected_connections)
-              Future.resolved
+          f = connect_to_all_peers(seed_connections, keyspace)
+          f.flat_map do |all_connections|
+            new_connections = all_connections - seed_connections
+            if new_connections.size > 0
+              f = use_keyspace(new_connections, keyspace)
+              f.on_value do
+                @connection_manager.add_connections(new_connections)
+              end
+              f
             elsif remaning_attempts > 0
               timeout = 2**(MAX_RECONNECTION_ATTEMPTS - remaning_attempts)
               @logger.debug('Scheduling new peer discovery in %ds' % timeout)

data/lib/cql/client/asynchronous_prepared_statement.rb

@@ -18,7 +18,7 @@ module Cql
         futures = connection_manager.map do |connection|
           statement.prepare(connection)
         end
-        Future.all(*futures).map { statement }
+        Future.all(*futures).map(statement)
       rescue => e
         Future.failed(e)
       end
@@ -36,6 +36,21 @@ module Cql
         Future.failed(e)
       end
 
+      def batch(type=:logged, options=nil)
+        if type.is_a?(Hash)
+          options = type
+          type = :logged
+        end
+        b = AsynchronousBatch.new(type, @execute_options_decoder, @connection_manager, options)
+        pb = AsynchronousPreparedStatementBatch.new(self, b)
+        if block_given?
+          yield pb
+          pb.execute
+        else
+          pb
+        end
+      end
+
       # @private
       def prepare(connection)
         prepare_request = Protocol::PrepareRequest.new(@cql)
@@ -46,25 +61,45 @@ module Cql
             # is that we assign the same data multiple times
             @raw_metadata = response.metadata
             @metadata = ResultMetadata.new(@raw_metadata)
+            @raw_result_metadata = response.result_metadata
+            if @raw_result_metadata
+              @result_metadata = ResultMetadata.new(@raw_result_metadata)
+            end
           end
           hex_id = response.id.each_byte.map { |x| x.to_s(16).rjust(2, '0') }.join('')
           @logger.debug('Statement %s prepared on node %s (%s:%d)' % [hex_id, connection[:host_id].to_s, connection.host, connection.port])
         end
-        f.map { self }
+        f.map(self)
+      end
+
+      # @private
+      def add_to_batch(batch, connection, bound_args)
+        statement_id = connection[self]
+        unless statement_id
+          raise NotPreparedError
+        end
+        unless bound_args.size == @raw_metadata.size
+          raise ArgumentError, "Expected #{@raw_metadata.size} arguments, got #{bound_args.size}"
+        end
+        batch.add_prepared(statement_id, @raw_metadata, bound_args)
       end
 
       private
 
       def run(args, connection)
-        statement_id = connection[self]
         bound_args = args.shift(@raw_metadata.size)
         unless bound_args.size == @raw_metadata.size && args.size <= 1
           raise ArgumentError, "Expected #{@raw_metadata.size} arguments, got #{bound_args.size}"
         end
-        consistency, timeout, trace = @execute_options_decoder.decode_options(args.last)
+        options = @execute_options_decoder.decode_options(args.last)
         statement_id = connection[self]
-        request = Protocol::ExecuteRequest.new(statement_id, @raw_metadata, bound_args, consistency, trace)
-        @request_runner.execute(connection, request, timeout)
+        request_metadata = @raw_result_metadata.nil?
+        request = Protocol::ExecuteRequest.new(statement_id, @raw_metadata, bound_args, request_metadata, options[:consistency], options[:serial_consistency], options[:page_size], options[:paging_state], options[:trace])
+        f = @request_runner.execute(connection, request, options[:timeout], @raw_result_metadata)
+        if options.include?(:page_size)
+          f = f.map { |result| AsynchronousPreparedPagedQueryResult.new(self, request, result, options) }
+        end
+        f
       end
     end
   end

data/lib/cql/client/authenticators.rb

@@ -0,0 +1,46 @@
+# encoding: utf-8
+
+module Cql
+  module Client
+    # Auth provider used for Cassandra's built in authentication.
+    #
+    # There is no need to create instances of this class to pass as `:auth_provider`
+    # to {Cql::Client.connect}, instead use the `:credentials` option and one
+    # will be created automatically for you.
+    class PlainTextAuthProvider
+      def initialize(username, password)
+        @username = username
+        @password = password
+      end
+
+      def create_authenticator(authentication_class)
+        if authentication_class == PASSWORD_AUTHENTICATOR_FQCN
+          PlainTextAuthenticator.new(@username, @password)
+        end
+      end
+
+      private
+
+      PASSWORD_AUTHENTICATOR_FQCN = 'org.apache.cassandra.auth.PasswordAuthenticator'.freeze
+    end
+
+    # Authenticator used for Cassandra's built in authentication,
+    # see {Cql::Client::PlainTextAuthProvider}
+    class PlainTextAuthenticator
+      def initialize(username, password)
+        @username = username
+        @password = password
+      end
+
+      def initial_response
+        "\x00#{@username}\x00#{@password}"
+      end
+
+      def challenge_response(token)
+      end
+
+      def authentication_successful(token)
+      end
+    end
+  end
+end

data/lib/cql/client/batch.rb

@@ -0,0 +1,115 @@
+# encoding: utf-8
+
+module Cql
+  module Client
+    # @private
+    class AsynchronousBatch < Batch
+      def initialize(type, execute_options_decoder, connection_manager, options=nil)
+        raise ArgumentError, "Unknown batch type: #{type}" unless BATCH_TYPES.include?(type)
+        @type = type
+        @execute_options_decoder = execute_options_decoder
+        @connection_manager = connection_manager
+        @options = options
+        @request_runner = RequestRunner.new
+        @parts = []
+      end
+
+      def add(*args)
+        @parts << args
+        nil
+      end
+
+      def execute(options=nil)
+        options = @execute_options_decoder.decode_options(@options, options)
+        connection = nil
+        attempts = 0
+        begin
+          connection = @connection_manager.random_connection
+          request = Protocol::BatchRequest.new(BATCH_TYPES[@type], options[:consistency], options[:trace])
+          @parts.each do |cql_or_statement, *bound_args|
+            if cql_or_statement.is_a?(String)
+              type_hints = nil
+              if bound_args.last.is_a?(Hash) && bound_args.last.include?(:type_hints)
+                bound_args = bound_args.dup
+                type_hints = bound_args.pop[:type_hints]
+              end
+              request.add_query(cql_or_statement, bound_args, type_hints)
+            else
+              cql_or_statement.add_to_batch(request, connection, bound_args)
+            end
+          end
+        rescue NotPreparedError
+          attempts += 1
+          if attempts < 3
+            retry
+          else
+            raise
+          end
+        end
+        @request_runner.execute(connection, request, options[:timeout])
+      end
+
+      private
+
+      BATCH_TYPES = {
+        :logged => Protocol::BatchRequest::LOGGED_TYPE,
+        :unlogged => Protocol::BatchRequest::UNLOGGED_TYPE,
+        :counter => Protocol::BatchRequest::COUNTER_TYPE,
+      }.freeze
+    end
+
+    # @private
+    class SynchronousBatch < Batch
+      include SynchronousBacktrace
+
+      def initialize(asynchronous_batch)
+        @asynchronous_batch = asynchronous_batch
+      end
+
+      def async
+        @asynchronous_batch
+      end
+
+      def add(*args)
+        @asynchronous_batch.add(*args)
+      end
+
+      def execute(options=nil)
+        synchronous_backtrace { @asynchronous_batch.execute(options).value }
+      end
+    end
+
+    # @private
+    class AsynchronousPreparedStatementBatch < PreparedStatementBatch
+      def initialize(prepared_statement, batch)
+        @prepared_statement = prepared_statement
+        @batch = batch
+      end
+
+      def add(*args)
+        @batch.add(@prepared_statement, *args)
+      end
+
+      def execute(options=nil)
+        @batch.execute(options)
+      end
+    end
+
+    # @private
+    class SynchronousPreparedStatementBatch < PreparedStatementBatch
+      include SynchronousBacktrace
+
+      def initialize(asynchronous_batch)
+        @asynchronous_batch = asynchronous_batch
+      end
+
+      def add(*args)
+        @asynchronous_batch.add(*args)
+      end
+
+      def execute(options=nil)
+        synchronous_backtrace { @asynchronous_batch.execute(options).value }
+      end
+    end
+  end
+end