jamesgolick-cassandra 0.8.2

data/lib/cassandra.rb

@@ -0,0 +1,22 @@
+require 'rubygems'
+require 'thrift_client'
+require 'json' unless defined?(JSON)
+require 'simple_uuid'
+here = File.expand_path(File.dirname(__FILE__))
+
+$LOAD_PATH << "#{here}/../vendor/gen-rb"
+require "#{here}/../vendor/gen-rb/cassandra"
+
+$LOAD_PATH << "#{here}"
+
+require 'cassandra/helpers'
+require 'cassandra/array'
+require 'cassandra/time'
+require 'cassandra/comparable'
+require 'cassandra/long'
+require 'cassandra/ordered_hash'
+require 'cassandra/columns'
+require 'cassandra/protocol'
+require 'cassandra/cassandra'
+require 'cassandra/constants'
+require 'cassandra/debug' if ENV['DEBUG']

data/lib/cassandra/array.rb

@@ -0,0 +1,8 @@
+
+class Array
+  def _flatten_once
+    result = []
+    each { |el| result.concat(Array(el)) }
+    result
+  end  
+end

data/lib/cassandra/cassandra.rb

@@ -0,0 +1,314 @@
+
+=begin rdoc
+Create a new Cassandra client instance. Accepts a keyspace name, and optional host and port.
+
+  client = Cassandra.new('twitter', '127.0.0.1:9160')
+
+You can then make calls to the server via the <tt>client</tt> instance.
+
+  client.insert(:UserRelationships, "5", {"user_timeline" => {SimpleUUID::UUID.new => "1"}})
+  client.get(:UserRelationships, "5", "user_timeline")
+
+For read methods, valid option parameters are:
+
+<tt>:count</tt>:: How many results to return. Defaults to 100.
+<tt>:start</tt>:: Column name token at which to start iterating, inclusive. Defaults to nil, which means the first column in the collation order.
+<tt>:finish</tt>:: Column name token at which to stop iterating, inclusive. Defaults to nil, which means no boundary.
+<tt>:reversed</tt>:: Swap the direction of the collation order.
+<tt>:consistency</tt>:: The consistency level of the request. Defaults to <tt>Cassandra::Consistency::ONE</tt> (one node must respond). Other valid options are <tt>Cassandra::Consistency::ZERO</tt>, <tt>Cassandra::Consistency::QUORUM</tt>, and <tt>Cassandra::Consistency::ALL</tt>.
+
+Note that some read options have no relevance in some contexts.
+
+For write methods, valid option parameters are:
+
+<tt>:timestamp </tt>:: The transaction timestamp. Defaults to the current time in milliseconds. This is used for conflict resolution by the server; you normally never need to change it.
+<tt>:consistency</tt>:: See above.
+
+For the initial client instantiation, you may also pass in <tt>:thrift_client<tt> with a ThriftClient subclass attached. On connection, that class will be used instead of the default ThriftClient class, allowing you to add additional behavior to the connection (e.g. query logging).
+
+=end rdoc
+
+class Cassandra
+  include Columns
+  include Protocol
+  include Helpers
+
+  class AccessError < StandardError #:nodoc:
+  end
+
+  module Consistency
+    include CassandraThrift::ConsistencyLevel
+  end
+
+  WRITE_DEFAULTS = {
+    :count => 1000,
+    :timestamp => nil,
+    :consistency => Consistency::ONE
+  }.freeze
+
+  READ_DEFAULTS = {
+    :count => 100,
+    :start => nil,
+    :finish => nil,
+    :reversed => false,
+    :consistency => Consistency::ONE
+  }.freeze
+  
+  THRIFT_DEFAULTS = {
+    :transport_wrapper => Thrift::BufferedTransport,
+    :thrift_client_class => ThriftClient
+  }.freeze
+
+  attr_reader :keyspace, :servers, :schema, :thrift_client_options, :thrift_client_class
+
+  # Create a new Cassandra instance and open the connection.
+  def initialize(keyspace, servers = "127.0.0.1:9160", thrift_client_options = {})
+    @is_super = {}
+    @column_name_class = {}
+    @sub_column_name_class = {}
+    @auto_discover_nodes = true
+    @thrift_client_options = THRIFT_DEFAULTS.merge(thrift_client_options)
+    @thrift_client_class = @thrift_client_options[:thrift_client_class]
+    @keyspace = keyspace
+    @servers = Array(servers)
+  end
+
+  def disable_node_auto_discovery!
+    @auto_discover_nodes = false
+  end
+
+  def disconnect!
+    @client.disconnect!
+    @client = nil
+  end
+
+  def keyspaces
+    @keyspaces ||= client.get_string_list_property("keyspaces")
+  end
+
+  def inspect
+    "#<Cassandra:#{object_id}, @keyspace=#{keyspace.inspect}, @schema={#{
+      schema(false).map {|name, hash| ":#{name} => #{hash['type'].inspect}"}.join(', ')
+    }}, @servers=#{servers.inspect}>"
+  end
+
+### Write
+
+  # Insert a row for a key. Pass a flat hash for a regular column family, and
+  # a nested hash for a super column family. Supports the <tt>:consistency</tt>
+  # and <tt>:timestamp</tt> options.
+  def insert(column_family, key, hash, options = {})
+    column_family, _, _, options = extract_and_validate_params(column_family, key, [options], WRITE_DEFAULTS)
+
+    timestamp = options[:timestamp] || Time.stamp
+    mutation_map = if is_super(column_family)
+      {
+        key => {
+          column_family => hash.collect{|k,v| _super_insert_mutation(column_family, k, v, timestamp) }
+        }
+      }
+    else
+      {
+        key => {
+          column_family => hash.collect{|k,v| _standard_insert_mutation(column_family, k, v, timestamp)}
+        }
+      }
+    end
+
+    @batch ? @batch << [mutation_map, options[:consistency]] : _mutate(mutation_map, options[:consistency])
+  end
+
+  
+  ## Delete
+
+  # _mutate the element at the column_family:key:[column]:[sub_column]
+  # path you request. Supports the <tt>:consistency</tt> and <tt>:timestamp</tt>
+  # options.
+  def remove(column_family, key, *columns_and_options)
+    column_family, column, sub_column, options = extract_and_validate_params(column_family, key, columns_and_options, WRITE_DEFAULTS)
+
+    args = {:column_family => column_family}
+    columns = is_super(column_family) ? {:super_column => column, :column => sub_column} : {:column => column}
+    column_path = CassandraThrift::ColumnPath.new(args.merge(columns))
+    
+    mutation = [:remove, [key, column_path, options[:timestamp] || Time.stamp, options[:consistency]]]
+    
+    @batch ? @batch << mutation : _remove(*mutation[1])
+  end
+
+  # Remove all rows in the column family you request. Supports options
+  # <tt>:consistency</tt> and <tt>:timestamp</tt>.
+  # FIXME May not currently delete all records without multiple calls. Waiting
+  # for ranged remove support in Cassandra.
+  def clear_column_family!(column_family, options = {})
+    each_key(column_family) do |key|
+      remove(column_family, key, options)
+    end
+  end
+
+  # Remove all rows in the keyspace. Supports options <tt>:consistency</tt> and
+  # <tt>:timestamp</tt>.
+  # FIXME May not currently delete all records without multiple calls. Waiting
+  # for ranged remove support in Cassandra.
+  def clear_keyspace!(options = {})
+    schema.keys.each { |column_family| clear_column_family!(column_family, options) }
+  end
+
+### Read
+
+  # Count the elements at the column_family:key:[super_column] path you
+  # request. Supports the <tt>:consistency</tt> option.
+  def count_columns(column_family, key, *columns_and_options)
+    column_family, super_column, _, options = 
+      extract_and_validate_params(column_family, key, columns_and_options, READ_DEFAULTS)      
+    _count_columns(column_family, key, super_column, options[:consistency])
+  end
+
+  # Multi-key version of Cassandra#count_columns. Supports options <tt>:count</tt>,
+  # <tt>:start</tt>, <tt>:finish</tt>, <tt>:reversed</tt>, and <tt>:consistency</tt>.
+  # FIXME Not real multi; needs server support
+  def multi_count_columns(column_family, keys, *options)
+    OrderedHash[*keys.map { |key| [key, count_columns(column_family, key, *options)] }._flatten_once]
+  end
+
+  # Return a list of single values for the elements at the
+  # column_family:key:column[s]:[sub_columns] path you request. Supports the
+  # <tt>:consistency</tt> option.
+  def get_columns(column_family, key, *columns_and_options)
+    column_family, columns, sub_columns, options = 
+      extract_and_validate_params(column_family, key, columns_and_options, READ_DEFAULTS)      
+    _get_columns(column_family, key, columns, sub_columns, options[:consistency])
+  end
+
+  # Multi-key version of Cassandra#get_columns. Supports the <tt>:consistency</tt>
+  # option.
+  # FIXME Not real multi; needs to use a Column predicate
+  def multi_get_columns(column_family, keys, *options)
+    OrderedHash[*keys.map { |key| [key, get_columns(column_family, key, *options)] }._flatten_once]
+  end
+
+  # Return a hash (actually, a Cassandra::OrderedHash) or a single value
+  # representing the element at the column_family:key:[column]:[sub_column]
+  # path you request. Supports options <tt>:count</tt>, <tt>:start</tt>,
+  # <tt>:finish</tt>, <tt>:reversed</tt>, and <tt>:consistency</tt>.
+  def get(column_family, key, *columns_and_options)
+    multi_get(column_family, [key], *columns_and_options)[key]
+  end
+
+  # Multi-key version of Cassandra#get. Supports options <tt>:count</tt>,
+  # <tt>:start</tt>, <tt>:finish</tt>, <tt>:reversed</tt>, and <tt>:consistency</tt>.
+  def multi_get(column_family, keys, *columns_and_options)
+    column_family, column, sub_column, options = 
+      extract_and_validate_params(column_family, keys, columns_and_options, READ_DEFAULTS)
+
+    hash = _multiget(column_family, keys, column, sub_column, options[:count], options[:start], options[:finish], options[:reversed], options[:consistency])
+    # Restore order
+    ordered_hash = OrderedHash.new
+    keys.each { |key| ordered_hash[key] = hash[key] || (OrderedHash.new if is_super(column_family) and !sub_column) }
+    ordered_hash
+  end
+
+  # Return true if the column_family:key:[column]:[sub_column] path you
+  # request exists. Supports the <tt>:consistency</tt> option.
+  def exists?(column_family, key, *columns_and_options)
+    column_family, column, sub_column, options = 
+      extract_and_validate_params(column_family, key, columns_and_options, READ_DEFAULTS)
+    if column
+      _multiget(column_family, [key], column, sub_column, 1, nil, nil, nil, options[:consistency])[key]
+    else
+      _multiget(column_family, [key], nil, nil, 1, '', '', false, options[:consistency])[key]
+    end
+  end
+
+  # Return a list of keys in the column_family you request. Requires the
+  # table to be partitioned with OrderPreservingHash. Supports the
+  # <tt>:count</tt>, <tt>:start</tt>, <tt>:finish</tt>, and <tt>:consistency</tt>
+  # options.
+  def get_range(column_family, options = {})
+    column_family, _, _, options = 
+      extract_and_validate_params(column_family, "", [options], READ_DEFAULTS)
+    _get_range(column_family, options[:start].to_s, options[:finish].to_s, options[:count], options[:consistency])
+  end
+
+  # Count all rows in the column_family you request. Requires the table
+  # to be partitioned with OrderPreservingHash. Supports the <tt>:start</tt>,
+  # <tt>:finish</tt>, and <tt>:consistency</tt> options.
+  def count_range(column_family, options = {})
+    get_range(column_family, options).select{|r| r.columns.length > 0}.compact.length
+  end
+
+  # Open a batch operation and yield. Inserts and deletes will be queued until
+  # the block closes, and then sent atomically to the server.  Supports the
+  # <tt>:consistency</tt> option, which overrides the consistency set in
+  # the individual commands.
+  def batch(options = {})
+    _, _, _, options = 
+      extract_and_validate_params(schema.keys.first, "", [options], WRITE_DEFAULTS)
+
+    @batch = []
+    yield
+    compact_mutations!
+
+    @batch.each do |mutation|
+      case mutation.first
+      when :remove
+        _remove(*mutation[1])
+      else
+        _mutate(*mutation)
+      end
+    end
+  ensure
+    @batch = nil
+  end
+
+  protected
+
+  def calling_method
+    "#{self.class}##{caller[0].split('`').last[0..-3]}"
+  end
+
+  # Roll up queued mutations, to improve atomicity.
+  def compact_mutations!
+    #TODO re-do this rollup
+  end
+
+  def schema(load=true)
+    if !load && !@schema
+      []
+    else
+      @schema ||= client.describe_keyspace(@keyspace)
+    end
+  end
+
+  def client
+    reconnect! if @client.nil?
+    @client
+  end
+
+  def reconnect!
+    @servers = all_nodes
+    @client = new_client
+    check_keyspace
+  end
+
+  def check_keyspace
+    unless (keyspaces = client.get_string_list_property("keyspaces")).include?(@keyspace)
+      raise AccessError, "Keyspace #{@keyspace.inspect} not found. Available: #{keyspaces.inspect}"
+    end
+  end
+
+  def new_client
+    thrift_client_class.new(CassandraThrift::Cassandra::Client, @servers, @thrift_client_options)
+  end
+
+  def all_nodes
+    if @auto_discover_nodes
+      ips = ::JSON.parse(new_client.get_string_property('token map')).values
+      port = @servers.first.split(':').last
+      ips.map{|ip| "#{ip}:#{port}" }
+    else
+      @servers
+    end
+  end
+
+end

data/lib/cassandra/columns.rb

@@ -0,0 +1,106 @@
+
+class Cassandra
+  # A bunch of crap, mostly related to introspecting on column types
+  module Columns #:nodoc:
+    private
+
+    def is_super(column_family)
+      @is_super[column_family] ||= column_family_property(column_family, 'Type') == "Super"
+    end
+
+    def column_name_class(column_family)
+      @column_name_class[column_family] ||= column_name_class_for_key(column_family, "CompareWith")
+    end
+
+    def sub_column_name_class(column_family)
+      @sub_column_name_class[column_family] ||= column_name_class_for_key(column_family, "CompareSubcolumnsWith")
+    end
+
+    def column_name_class_for_key(column_family, comparator_key)
+      property = column_family_property(column_family, comparator_key)
+      property =~ /.*\.(.*?)$/
+      case $1
+      when "LongType" then Long
+      when "LexicalUUIDType", "TimeUUIDType" then SimpleUUID::UUID
+      else 
+        String # UTF8, Ascii, Bytes, anything else
+      end
+    end
+
+    def column_family_property(column_family, key)
+      unless schema[column_family]
+        raise AccessError, "Invalid column family \"#{column_family}\""
+      end
+      schema[column_family][key]
+    end
+
+    def multi_column_to_hash!(hash)
+      hash.each do |key, column_or_supercolumn|
+        hash[key] = (column_or_supercolumn.column.value if column_or_supercolumn.column)
+      end
+    end
+
+    def multi_columns_to_hash!(column_family, hash)
+      hash.each do |key, columns| 
+        hash[key] = columns_to_hash(column_family, columns)
+      end
+    end
+
+    def multi_sub_columns_to_hash!(column_family, hash)
+      hash.each do |key, sub_columns| 
+        hash[key] = sub_columns_to_hash(column_family, sub_columns)
+      end
+    end
+
+    def columns_to_hash(column_family, columns)
+      columns_to_hash_for_classes(columns, column_name_class(column_family), sub_column_name_class(column_family))
+    end
+
+    def sub_columns_to_hash(column_family, columns)
+      columns_to_hash_for_classes(columns, sub_column_name_class(column_family))
+    end
+
+    def columns_to_hash_for_classes(columns, column_name_class, sub_column_name_class = nil)
+      hash = OrderedHash.new
+      Array(columns).each do |c|
+        c = c.super_column || c.column if c.is_a?(CassandraThrift::ColumnOrSuperColumn)
+        hash[column_name_class.new(c.name)] = case c
+        when CassandraThrift::SuperColumn            
+          columns_to_hash_for_classes(c.columns, sub_column_name_class) # Pop the class stack, and recurse
+        when CassandraThrift::Column
+          c.value
+        end
+      end
+      hash
+    end
+
+    def _standard_insert_mutation(column_family, column_name, value, timestamp)
+      CassandraThrift::Mutation.new(
+        :column_or_supercolumn => CassandraThrift::ColumnOrSuperColumn.new(
+          :column => CassandraThrift::Column.new(
+            :name      => column_name_class(column_family).new(column_name).to_s,
+            :value     => value,
+            :timestamp => timestamp
+          )
+        )
+      )
+    end
+
+    def _super_insert_mutation(column_family, super_column_name, sub_columns, timestamp)
+      CassandraThrift::Mutation.new(:column_or_supercolumn => 
+        CassandraThrift::ColumnOrSuperColumn.new(
+          :super_column => CassandraThrift::SuperColumn.new(
+            :name => column_name_class(column_family).new(super_column_name).to_s,
+            :columns => sub_columns.collect { |sub_column_name, sub_column_value|
+              CassandraThrift::Column.new(
+                :name      => sub_column_name_class(column_family).new(sub_column_name).to_s,
+                :value     => sub_column_value.to_s,
+                :timestamp => timestamp
+              )
+            }
+          )
+        )
+      )
+    end
+  end
+end