cql-rb 1.2.2 → 2.0.0.pre0

data/lib/cql/client/connector.rb

@@ -0,0 +1,255 @@
+# encoding: utf-8
+
+module Cql
+  module Client
+    # @private
+    class ClusterConnector
+      def initialize(sequence, logger)
+        @sequence = sequence
+        @logger = logger
+      end
+
+      def connect_all(hosts, connections_per_node)
+        connections = hosts.flat_map do |host|
+          Array.new(connections_per_node) do
+            f = @sequence.connect(host)
+            f.on_value { |connection| register_logging(connection) }
+            f.recover do |error|
+              @logger.warn('Failed connecting to node at %s: %s' % [host, error.message])
+              FailedConnection.new(error, host)
+            end
+          end
+        end
+        Future.all(*connections).map do |connections|
+          connected_connections = connections.select(&:connected?)
+          if connected_connections.empty?
+            e = connections.first.error
+            if e.is_a?(Cql::QueryError) && e.code == 0x100
+              e = AuthenticationError.new(e.message)
+            end
+            raise e
+          end
+          connected_connections
+        end
+      end
+
+      private
+
+      def register_logging(connection)
+        args = [connection[:host_id], connection.host, connection.port, connection[:data_center]]
+        @logger.info('Connected to node %s at %s:%d in data center %s' % args)
+        connection.on_closed do |cause|
+          message = 'Connection to node %s at %s:%d in data center %s closed' % args
+          if cause
+            message << (' unexpectedly: %s' % cause.message)
+            @logger.warn(message)
+          else
+            @logger.info(message)
+          end
+        end
+      end
+    end
+
+    # @private
+    class Connector
+      def initialize(steps)
+        @steps = steps.dup
+      end
+
+      def connect(host)
+        pending_connection = PendingConnection.new(host)
+        seed = Future.resolved(pending_connection)
+        f = @steps.reduce(seed) do |chain, step|
+          chain.flat_map do |pending_connection|
+            step.run(pending_connection)
+          end
+        end
+        f.map do |pending_connection|
+          pending_connection.connection
+        end
+      end
+    end
+
+    # @private
+    class ConnectStep
+      def initialize(io_reactor, port, connection_timeout, logger)
+        @io_reactor = io_reactor
+        @port = port
+        @connection_timeout = connection_timeout
+        @logger = logger
+      end
+
+      def run(pending_connection)
+        @logger.debug('Connecting to node at %s:%d' % [pending_connection.host, @port])
+        @io_reactor.connect(pending_connection.host, @port, @connection_timeout).map do |connection|
+          pending_connection.with_connection(connection)
+        end
+      end
+    end
+
+    # @private
+    class CacheOptionsStep
+      def run(pending_connection)
+        f = pending_connection.execute(Protocol::OptionsRequest.new)
+        f.on_value do |supported_options|
+          pending_connection[:cql_version] = supported_options['CQL_VERSION']
+          pending_connection[:compression] = supported_options['COMPRESSION']
+        end
+        f.map(pending_connection)
+      end
+    end
+
+    # @private
+    class InitializeStep
+      def initialize(cql_version, compressor, logger)
+        @cql_version = cql_version
+        @compressor = compressor
+        @logger = logger
+      end
+
+      def run(pending_connection)
+        compression = @compressor && @compressor.algorithm
+        supported_algorithms = pending_connection[:compression]
+        if @compressor && !supported_algorithms.include?(@compressor.algorithm)
+          @logger.warn(%[Compression algorithm "#{@compressor.algorithm}" not supported (server supports "#{supported_algorithms.join('", "')}")])
+          compression = nil
+        elsif @compressor
+          @logger.debug('Using "%s" compression' % @compressor.algorithm)
+        end
+        f = pending_connection.execute(Protocol::StartupRequest.new(@cql_version, compression))
+        f.map do |startup_response|
+          if startup_response.is_a?(AuthenticationRequired)
+            pending_connection.with_authentication_class(startup_response.authentication_class)
+          else
+            pending_connection
+          end
+        end
+      end
+    end
+
+    # @private
+    class SaslAuthenticationStep
+      def initialize(auth_provider)
+        @auth_provider = auth_provider
+      end
+
+      def run(pending_connection)
+        if pending_connection.authentication_class
+          begin
+            authenticator = @auth_provider && @auth_provider.create_authenticator(pending_connection.authentication_class)
+            if authenticator
+              token = authenticator.initial_response
+              challenge_cycle(pending_connection, authenticator, token)
+            elsif @auth_provider
+              Future.failed(AuthenticationError.new('Auth provider does not support the required authentication class "%s" and/or protocol version %d' % [pending_connection.authentication_class, @protocol_version]))
+            else
+              Future.failed(AuthenticationError.new('Server requested authentication, but no auth provider found'))
+            end
+          rescue => e
+            Future.failed(AuthenticationError.new('Auth provider raised an error: %s' % e.message))
+          end
+        else
+          Future.resolved(pending_connection)
+        end
+      end
+
+      def challenge_cycle(pending_connection, authenticator, response_token)
+        request = Protocol::AuthResponseRequest.new(response_token)
+        f = pending_connection.execute(request) { |raw_response| raw_response }
+        f.flat_map do |response|
+          case response
+          when Protocol::AuthChallengeResponse
+            token = authenticator.challenge_response(response.token)
+            challenge_cycle(pending_connection, authenticator, token)
+          when Protocol::AuthSuccessResponse
+            authenticator.authentication_successful(response.token)
+            Future.resolved(pending_connection)
+          else
+            Future.resolved(pending_connection)
+          end
+        end
+      end
+    end
+
+    # @private
+    class CredentialsAuthenticationStep
+      def initialize(credentials)
+        @credentials = credentials
+      end
+
+      def run(pending_connection)
+        if pending_connection.authentication_class
+          if @credentials
+            request = Protocol::CredentialsRequest.new(@credentials)
+            pending_connection.execute(request).map(pending_connection)
+          else
+            Future.failed(AuthenticationError.new('Server requested authentication, but no credentials provided'))
+          end
+        else
+          Future.resolved(pending_connection)
+        end
+      end
+    end
+
+    # @private
+    class CachePropertiesStep
+      def run(pending_connection)
+        request = Protocol::QueryRequest.new('SELECT data_center, host_id FROM system.local', nil, nil, :one)
+        f = pending_connection.execute(request)
+        f.on_value do |result|
+          unless result.empty?
+            pending_connection[:host_id] = result.first['host_id']
+            pending_connection[:data_center] = result.first['data_center']
+          end
+        end
+        f.map(pending_connection)
+      end
+    end
+
+    # @private
+    class PendingConnection
+      attr_reader :host, :connection, :authentication_class
+
+      def initialize(host, connection=nil, authentication_class=nil)
+        @host = host
+        @connection = connection
+        @authentication_class = authentication_class
+        @request_runner = RequestRunner.new
+      end
+
+      def with_connection(connection)
+        self.class.new(host, connection, @authentication_class)
+      end
+
+      def with_authentication_class(authentication_class)
+        self.class.new(host, @connection, authentication_class)
+      end
+
+      def [](key)
+        @connection[key]
+      end
+
+      def []=(key, value)
+        @connection[key] = value
+      end
+
+      def execute(request, &block)
+        @request_runner.execute(@connection, request, nil, nil, &block)
+      end
+    end
+
+    # @private
+    class FailedConnection
+      attr_reader :error, :host
+
+      def initialize(error, host)
+        @error = error
+        @host = host
+      end
+
+      def connected?
+        false
+      end
+    end
+  end
+end

data/lib/cql/client/execute_options_decoder.rb

@@ -6,21 +6,37 @@ module Cql
     class ExecuteOptionsDecoder
       def initialize(default_consistency)
         @default_consistency = default_consistency
+        @default_options = {:consistency => @default_consistency}.freeze
       end
 
-      def decode_options(options_or_consistency)
-        consistency = @default_consistency
-        timeout = nil
-        trace = false
+      def decode_options(*args)
+        if args.empty?
+          @default_options
+        elsif args.size == 1
+          decode_one(args.first)
+        else
+          args.each_with_object({}) do |options_or_consistency, result|
+            result.merge!(decode_one(options_or_consistency))
+          end
+        end
+      end
+
+      private
+
+      def decode_one(options_or_consistency)
+        return @default_options unless options_or_consistency
         case options_or_consistency
         when Symbol
-          consistency = options_or_consistency
+          {:consistency => options_or_consistency}
         when Hash
-          consistency = options_or_consistency[:consistency] || consistency
-          timeout = options_or_consistency[:timeout]
-          trace = options_or_consistency[:trace]
+          if options_or_consistency.include?(:consistency)
+            options_or_consistency
+          else
+            options = options_or_consistency.dup
+            options[:consistency] = @default_consistency
+            options
+          end
         end
-        return consistency, timeout, trace
       end
     end
   end

data/lib/cql/client/keyspace_changer.rb

@@ -6,15 +6,15 @@ module Cql
     class KeyspaceChanger
       KEYSPACE_NAME_PATTERN = /^\w[\w\d_]*$|^"\w[\w\d_]*"$/
       
-      def initialize
-        @request_runner = RequestRunner.new
+      def initialize(request_runner=RequestRunner.new)
+        @request_runner = request_runner
       end
 
-      def use_keyspace(keyspace, connection)
+      def use_keyspace(connection, keyspace)
         return Future.resolved(connection) unless keyspace
         return Future.failed(InvalidKeyspaceNameError.new(%("#{keyspace}" is not a valid keyspace name))) unless valid_keyspace_name?(keyspace)
-        request = Protocol::QueryRequest.new("USE #{keyspace}", :one)
-        @request_runner.execute(connection, request).map { connection }
+        request = Protocol::QueryRequest.new("USE #{keyspace}", nil, nil, :one)
+        @request_runner.execute(connection, request).map(connection)
       end
 
       private

data/lib/cql/client/peer_discovery.rb

@@ -0,0 +1,33 @@
+# encoding: utf-8
+
+module Cql
+  module Client
+    class PeerDiscovery
+      def initialize(seed_connections)
+        @seed_connections = seed_connections
+        @connection = seed_connections.sample
+        @request_runner = RequestRunner.new
+      end
+
+      def new_hosts
+        request = Protocol::QueryRequest.new('SELECT peer, data_center, host_id, rpc_address FROM system.peers', nil, nil, :one)
+        response = @request_runner.execute(@connection, request)
+        response.map do |result|
+          result.each_with_object([]) do |row, new_peers|
+            if include?(row['host_id'], row['data_center'])
+              rpc_address = row['rpc_address'].to_s
+              rpc_address = row['peer'].to_s if rpc_address == '0.0.0.0'
+              new_peers << rpc_address
+            end
+          end
+        end
+      end
+
+      private
+
+      def include?(host_id, dc)
+        @seed_connections.any? { |c| c[:data_center] == dc } && @seed_connections.none? { |c| c[:host_id] == host_id }
+      end
+    end
+  end
+end

data/lib/cql/client/query_result.rb

@@ -14,10 +14,14 @@ module Cql
       attr_reader :trace_id
 
       # @private
-      def initialize(metadata, rows, trace_id)
+      attr_reader :paging_state
+
+      # @private
+      def initialize(metadata, rows, trace_id, paging_state)
         @metadata = ResultMetadata.new(metadata)
         @rows = rows
         @trace_id = trace_id
+        @paging_state = paging_state
       end
 
       # Returns whether or not there are any rows in this result set
@@ -34,5 +38,124 @@ module Cql
       end
       alias_method :each_row, :each
     end
+
+    class PagedQueryResult < QueryResult
+      def metadata
+        @result.metadata
+      end
+
+      def trace_id
+        @result.trace_id
+      end
+
+      def paging_state
+        @result.paging_state
+      end
+
+      def empty?
+        @result.empty?
+      end
+
+      def each(&block)
+        @result.each(&block)
+      end
+      alias_method :each_row, :each
+
+      def last_page?
+        @last_page
+      end
+
+      def next_page
+      end
+    end
+
+    # @private
+    class AsynchronousPagedQueryResult < PagedQueryResult
+      def initialize(request, result, options)
+        @request = request
+        @result = result
+        @result = result
+        @options = options.merge(paging_state: result.paging_state)
+        @last_page = !result.paging_state
+      end
+    end
+
+    # @private
+    class AsynchronousQueryPagedQueryResult < AsynchronousPagedQueryResult
+      def initialize(client, request, result, options)
+        super(request, result, options)
+        @client = client
+      end
+
+      def next_page
+        @client.execute(@request.cql, *@request.values, @options)
+      end
+    end
+
+    # @private
+    class AsynchronousPreparedPagedQueryResult < AsynchronousPagedQueryResult
+      def initialize(prepared_statement, request, result, options)
+        super(request, result, options)
+        @prepared_statement = prepared_statement
+      end
+
+      def next_page
+        @prepared_statement.execute(*@request.values, @options)
+      end
+    end
+
+    # @private
+    class SynchronousPagedQueryResult < PagedQueryResult
+      include SynchronousBacktrace
+
+      def initialize(asynchronous_result)
+        @result = asynchronous_result
+      end
+
+      def async
+        @result
+      end
+
+      def last_page?
+        @result.last_page?
+      end
+
+      def next_page
+        synchronous_backtrace { self.class.new(@result.next_page.value) }
+      end
+    end
+
+    # @private
+    class LazyQueryResult < QueryResult
+      def initialize(metadata, lazy_rows, trace_id, paging_state)
+        super(metadata, nil, trace_id, paging_state)
+        @raw_metadata = metadata
+        @lazy_rows = lazy_rows
+        @lock = Mutex.new
+      end
+
+      def empty?
+        ensure_materialized
+        super
+      end
+
+      def each(&block)
+        ensure_materialized
+        super
+      end
+      alias_method :each_row, :each
+
+      private
+
+      def ensure_materialized
+        unless @rows
+          @lock.synchronize do
+            unless @rows
+              @rows = @lazy_rows.materialize(@raw_metadata)
+            end
+          end
+        end
+      end
+    end
   end
 end