better-riak-client 1.0.5

data/erl_src/riak_search_test_backend.erl

@@ -0,0 +1,175 @@
+%% -------------------------------------------------------------------
+%%
+%% Copyright (c) 2007-2010 Basho Technologies, Inc.  All Rights Reserved.
+%%
+%% -------------------------------------------------------------------
+
+-module(riak_search_test_backend).
+-behavior(riak_search_backend).
+
+-export([
+         reset/0,
+         start/2,
+         stop/1,
+         index/2,
+         delete/2,
+         stream/6,
+         range/8,
+         info/5,
+         fold/3,
+         is_empty/1,
+         drop/1
+        ]).
+-export([
+         stream_results/3
+        ]).
+
+-include_lib("riak_search/include/riak_search.hrl").
+-define(T(P), list_to_atom("rs" ++ integer_to_list(P))).
+-record(state, {partition, table}).
+
+reset() ->
+    {ok, Ring} = riak_core_ring_manager:get_my_ring(),
+    [ catch ets:delete_all_objects(?T(P)) ||
+        P <- riak_core_ring:my_indices(Ring) ],
+    riak_search_config:clear(),
+    ok.
+
+start(Partition, _Config) ->
+    Table = ets:new(?T(Partition),
+                    [named_table, public, ordered_set]),
+    {ok, #state{partition=Partition, table=Table}}.
+
+stop(State) ->
+    maybe_delete(State).
+
+index(IFTVPKList, #state{table=Table}=State) ->
+    lists:foreach(
+      fun({I, F, T, V, P, K}) ->
+              Key = {b(I), b(F), b(T), b(V)},
+              case ets:lookup(Table, Key) of
+                  [{_, _, ExistingKeyClock}] ->
+                      if ExistingKeyClock > K ->
+                              %% stored data is newer
+                              ok;
+                         true ->
+                              %% stored data is older
+                              ets:update_element(Table, Key,
+                                                 [{2, P},{3, K}])
+                      end;
+                  [] ->
+                      ets:insert(Table, {Key, P, K})
+              end
+      end,
+      IFTVPKList),
+    {reply, {indexed, node()}, State}.
+
+delete(IFTVKList, State) ->
+    Table = State#state.table,
+    lists:foreach(fun(IFTVK) -> delete_fun(IFTVK, Table) end, IFTVKList),
+    {reply, {deleted, node()}, State}.
+
+delete_fun({I, F, T, V, K}, Table) ->
+    Key = {b(I), b(F), b(T), b(V)},
+    case ets:lookup(Table, Key) of
+        [{Key, _Props, ExistingKeyClock}] ->
+            if ExistingKeyClock > K ->
+                    %% stored data is newer
+                    ok;
+               true ->
+                    %% stored data is older
+                    ets:delete(Table, Key)
+            end;
+        [] ->
+            ok
+    end;
+delete_fun({I, F, T, V, _P, K}, Table) ->
+    %% copied idea from merge_index_backend
+    %% other operations include Props, though delete shouldn't
+    delete_fun({I, F, T, V, K}, Table).
+
+info(Index, Field, Term, Sender, State) ->
+    Count = ets:select_count(State#state.table,
+                             [{{{b(Index), b(Field), b(Term), '_'},
+                                '_', '_'},
+                               [],[true]}]),
+    riak_search_backend:info_response(Sender, [{Term, node(), Count}]),
+    noreply.
+
+-define(STREAM_SIZE, 100).
+
+range(Index, Field, StartTerm, EndTerm, _Size, FilterFun, Sender, State) ->
+    ST = b(StartTerm),
+    ET = b(EndTerm),
+    spawn(riak_search_ets_backend, stream_results,
+          [Sender,
+           FilterFun,
+           ets:select(State#state.table,
+                      [{{{b(Index), b(Field), '$1', '$2'}, '$3', '_'},
+                        [{'>=', '$1', ST}, {'=<', '$1', ET}],
+                        [{{'$2', '$3'}}]}],
+                      ?STREAM_SIZE)]),
+    noreply.
+
+stream(Index, Field, Term, FilterFun, Sender, State) ->
+    spawn(riak_search_ets_backend, stream_results,
+          [Sender,
+           FilterFun,
+           ets:select(State#state.table,
+                      [{{{b(Index), b(Field), b(Term), '$1'}, '$2', '_'},
+                        [], [{{'$1', '$2'}}]}],
+                      ?STREAM_SIZE)]),
+    noreply.
+
+stream_results(Sender, FilterFun, {Results0, Continuation}) ->
+    case lists:filter(fun({V,P}) -> FilterFun(V, P) end, Results0) of
+        [] ->
+            ok;
+        Results ->
+            riak_search_backend:response_results(Sender, Results)
+    end,
+    stream_results(Sender, FilterFun, ets:select(Continuation));
+stream_results(Sender, _, '$end_of_table') ->
+    riak_search_backend:response_done(Sender).
+
+fold(FoldFun, Acc, State) ->
+    Fun = fun({{I,F,T,V},P,K}, {OuterAcc, {{I,{F,T}},InnerAcc}}) ->
+                  %% same IFT, just accumulate doc/props/clock
+                  {OuterAcc, {{I,{F,T}},[{V,P,K}|InnerAcc]}};
+             ({{I,F,T,V},P,K}, {OuterAcc, {FoldKey, VPKList}}) ->
+                  %% finished a string of IFT, send it off
+                  %% (sorted order is assumed)
+                  NewOuterAcc = FoldFun(FoldKey, VPKList, OuterAcc),
+                  {NewOuterAcc, {{I,{F,T}},[{V,P,K}]}};
+             ({{I,F,T,V},P,K}, {OuterAcc, undefined}) ->
+                  %% first round through the fold - just start building
+                  {OuterAcc, {{I,{F,T}},[{V,P,K}]}}
+          end,
+    {OuterAcc0, Final} = ets:foldl(Fun, {Acc, undefined}, State#state.table),
+    OuterAcc = case Final of
+                   {FoldKey, VPKList} ->
+                       %% one last IFT to send off
+                       FoldFun(FoldKey, VPKList, OuterAcc0);
+                   undefined ->
+                       %% this partition was empty
+                       OuterAcc0
+               end,
+    {reply, OuterAcc, State}.
+
+is_empty(State) ->
+    0 == ets:info(State#state.table, size).
+
+drop(State) ->
+    maybe_delete(State).
+
+maybe_delete(State) ->
+    case lists:member(State#state.table, ets:all()) of
+        true ->
+            ets:delete(State#state.table),
+            ok;
+        false ->
+            ok
+    end.
+
+b(Binary) when is_binary(Binary) -> Binary;
+b(List) when is_list(List) -> iolist_to_binary(List).

data/lib/riak/bucket.rb

@@ -0,0 +1,221 @@
+require 'riak/util/translation'
+require 'riak/client'
+require 'riak/robject'
+require 'riak/failed_request'
+
+module Riak
+  # Represents and encapsulates operations on a Riak bucket.  You may retrieve a bucket
+  # using {Client#bucket}, or create it manually and retrieve its meta-information later.
+  class Bucket
+    include Util::Translation
+
+    # (Riak Search) The precommit specification for kv/search integration
+    SEARCH_PRECOMMIT_HOOK = {"mod" => "riak_search_kv_hook", "fun" => "precommit"}
+
+    # @return [Riak::Client] the associated client
+    attr_reader :client
+
+    # @return [String] the bucket name
+    attr_reader :name
+
+    # Create a Riak bucket manually.
+    # @param [Client] client the {Riak::Client} for this bucket
+    # @param [String] name the name of the bucket
+    def initialize(client, name)
+      raise ArgumentError, t("client_type", :client => client.inspect) unless Client === client
+      raise ArgumentError, t("string_type", :string => name.inspect) unless String === name
+      @client, @name = client, name
+    end
+
+    # Retrieves a list of keys in this bucket.
+    # If a block is given, keys will be streamed through
+    # the block (useful for large buckets). When streaming,
+    # results of the operation will not be returned to the caller.
+    # @yield [Array<String>] a list of keys from the current chunk
+    # @return [Array<String>] Keys in this bucket
+    # @note This operation has serious performance implications and
+    #    should not be used in production applications.
+    def keys(&block)
+      warn(t('list_keys', :backtrace => caller.join("\n    "))) unless Riak.disable_list_keys_warnings
+      if block_given?
+        @client.list_keys(self, &block)
+      else
+        @client.list_keys(self)
+      end
+    end
+
+    # Sets internal properties on the bucket
+    # Note: this results in a request to the Riak server!
+    # @param [Hash] properties new properties for the bucket
+    # @option properties [Fixnum] :n_val (3) The N value (replication factor)
+    # @option properties [true,false] :allow_mult (false) Whether to permit object siblings
+    # @option properties [true,false] :last_write_wins (false) Whether to ignore vclocks
+    # @option properties [Array<Hash>] :precommit ([]) precommit hooks
+    # @option properties [Array<Hash>] :postcommit ([])postcommit hooks
+    # @option properties [Fixnum,String] :r ("quorum") read quorum (numeric or
+    # symbolic)
+    # @option properties [Fixnum,String] :w ("quorum") write quorum (numeric or
+    # symbolic)
+    # @option properties [Fixnum,String] :dw ("quorum") durable write quorum
+    # (numeric or symbolic)
+    # @option properties [Fixnum,String] :rw ("quorum") delete quorum (numeric or
+    # symbolic)
+    # @return [Hash] the merged bucket properties
+    # @raise [FailedRequest] if the new properties were not accepted by the Riakserver
+    # @see #n_value, #allow_mult, #r, #w, #dw, #rw
+    def props=(properties)
+      raise ArgumentError, t("hash_type", :hash => properties.inspect) unless Hash === properties
+      props.merge!(properties)
+      @client.set_bucket_props(self, properties)
+      props
+    end
+    alias :'properties=' :'props='
+
+    # @return [Hash] Internal Riak bucket properties.
+    # @see #props=
+    def props
+      @props ||= @client.get_bucket_props(self)
+    end
+    alias :properties :props
+
+    # Retrieve an object from within the bucket.
+    # @param [String] key the key of the object to retrieve
+    # @param [Hash] options query parameters for the request
+    # @option options [Fixnum] :r - the read quorum for the request - how many nodes should concur on the read
+    # @return [Riak::RObject] the object
+    # @raise [FailedRequest] if the object is not found or some other error occurs
+    def get(key, options={})
+      @client.get_object(self, key, options)
+    end
+    alias :[] :get
+
+    # Create a new blank object
+    # @param [String] key the key of the new object
+    # @return [RObject] the new, unsaved object
+    def new(key=nil)
+      RObject.new(self, key).tap do |obj|
+        obj.content_type = "application/json"
+      end
+    end
+
+    # Fetches an object if it exists, otherwise creates a new one with the given key
+    # @param [String] key the key to fetch or create
+    # @return [RObject] the new or existing object
+    def get_or_new(key, options={})
+      begin
+        get(key, options)
+      rescue Riak::FailedRequest => fr
+        if fr.not_found?
+          new(key)
+        else
+          raise fr
+        end
+      end
+    end
+
+    # Checks whether an object exists in Riak.
+    # @param [String] key the key to check
+    # @param [Hash] options quorum options
+    # @option options [Fixnum] :r - the read quorum value for the request (R)
+    # @return [true, false] whether the key exists in this bucket
+    def exists?(key, options={})
+      begin
+        get(key, options)
+        true
+      rescue Riak::FailedRequest
+        false
+      end
+    end
+    alias :exist? :exists?
+
+    # Deletes a key from the bucket
+    # @param [String] key the key to delete
+    # @param [Hash] options quorum options
+    # @option options [Fixnum] :rw - the read/write quorum for the
+    #   delete
+    # @option options [String] :vclock - the vector clock of the
+    #   object being deleted
+    def delete(key, options={})
+      client.delete_object(self, key, options)
+    end
+
+    # Queries a secondary index on the bucket.
+    # @note This will only work if your Riak installation supports 2I.
+    # @param [String] index the name of the index
+    # @param [String,Integer,Range] query the value of the index, or a
+    #   Range of values to query
+    # @return [Array<String>] a list of keys that match the index
+    #   query
+    def get_index(index, query)
+      client.get_index(self, index, query)
+    end
+
+    # @return [true, false] whether the bucket allows divergent siblings
+    def allow_mult
+      props['allow_mult']
+    end
+
+    # Set the allow_mult property.  *NOTE* This will result in a PUT request to Riak.
+    # @param [true, false] value whether the bucket should allow siblings
+    def allow_mult=(value)
+      self.props = {'allow_mult' => value}
+      value
+    end
+
+    # @return [Fixnum] the N value, or number of replicas for this bucket
+    def n_value
+      props['n_val']
+    end
+    alias :n_val :n_value
+
+    # Set the N value (number of replicas). *NOTE* This will result in a PUT request to Riak.
+    # Setting this value after the bucket has objects stored in it may have unpredictable results.
+    # @param [Fixnum] value the number of replicas the bucket should keep of each object
+    def n_value=(value)
+      self.props = {'n_val' => value}
+      value
+    end
+    alias :'n_val=' :'n_value='
+
+    %w(r w dw rw).each do |q|
+      define_method(q) { props[q] }
+      define_method("#{q}=") { |value|
+        self.props = { q => value }
+        value
+      }
+    end
+
+    # (Riak Search) Installs a precommit hook that automatically indexes objects
+    # into riak_search.
+    def enable_index!
+      unless is_indexed?
+        self.props = {"precommit" => (props['precommit'] + [SEARCH_PRECOMMIT_HOOK]), "search" => true}
+      end
+    end
+
+    # (Riak Search) Removes the precommit hook that automatically indexes objects
+    # into riak_search.
+    def disable_index!
+      if is_indexed?
+        self.props = {"precommit" => (props['precommit'] - [SEARCH_PRECOMMIT_HOOK]), "search" => false}
+      end
+    end
+
+    # (Riak Search) Detects whether the bucket is automatically indexed into
+    # riak_search.
+    # @return [true,false] whether the bucket includes the search indexing hook
+    def is_indexed?
+      props['search'] == true || props['precommit'].include?(SEARCH_PRECOMMIT_HOOK)
+    end
+
+    # @return [String] a representation suitable for IRB and debugging output
+    def inspect
+      "#<Riak::Bucket {#{name}}#{" keys=[#{keys.join(',')}]" if defined?(@keys)}>"
+    end
+
+    # @return [true,false] whether the other is equivalent
+    def ==(other)
+      Bucket === other && other.client == client && other.name == name
+    end
+  end
+end

data/lib/riak/client/beefcake/messages.rb

@@ -0,0 +1,213 @@
+require 'beefcake'
+
+module Riak
+  class Client
+    # @private
+    class BeefcakeProtobuffsBackend
+      # Embedded messages
+      class RpbPair
+        include Beefcake::Message
+        required :key,   :bytes, 1
+        optional :value, :bytes, 2
+      end
+
+      class RpbBucketProps
+        include Beefcake::Message
+        optional :n_val,      :uint32, 1
+        optional :allow_mult, :bool,   2
+      end
+
+      class RpbLink
+        include Beefcake::Message
+        optional :bucket, :bytes, 1
+        optional :key,    :bytes, 2
+        optional :tag,    :bytes, 3
+      end
+
+      class RpbContent
+        include Beefcake::Message
+        required :value,            :bytes,  1
+        optional :content_type,     :bytes,  2
+        optional :charset,          :bytes,  3
+        optional :content_encoding, :bytes,  4
+        optional :vtag,             :bytes,  5
+        repeated :links,            RpbLink, 6
+        optional :last_mod,         :uint32, 7
+        optional :last_mod_usecs,   :uint32, 8
+        repeated :usermeta,         RpbPair, 9
+        repeated :indexes,          RpbPair, 10
+      end
+
+      # Primary messages
+      class RpbErrorResp
+        include Beefcake::Message
+        required :errmsg,  :bytes,  1
+        required :errcode, :uint32, 2
+      end
+
+      class RpbGetClientIdResp
+        include Beefcake::Message
+        required :client_id, :bytes, 1
+      end
+
+      class RpbSetClientIdReq
+        include Beefcake::Message
+        required :client_id, :bytes, 1
+      end
+
+      class RpbGetServerInfoResp
+        include Beefcake::Message
+        optional :node,           :bytes, 1
+        optional :server_version, :bytes, 2
+      end
+
+      class RpbGetReq
+        include Beefcake::Message
+        required :bucket,        :bytes,  1
+        required :key,           :bytes,  2
+        optional :r,             :uint32, 3
+        optional :pr,            :uint32, 4
+        optional :basic_quorum,  :bool,   5
+        optional :notfound_ok,   :bool,   6
+        optional :if_modified,   :bytes,  7
+        optional :head,          :bool,   8
+        optional :deletedvclock, :bool,   9
+      end
+
+      class RpbGetResp
+        include Beefcake::Message
+        repeated :content,   RpbContent, 1
+        optional :vclock,    :bytes,     2
+        optional :unchanged, :bool,      3
+      end
+
+      class RpbPutReq
+        include Beefcake::Message
+        required :bucket,          :bytes,     1
+        optional :key,             :bytes,     2
+        optional :vclock,          :bytes,     3
+        required :content,         RpbContent, 4
+        optional :w,               :uint32,    5
+        optional :dw,              :uint32,    6
+        optional :returnbody,      :bool,      7
+        optional :pw,              :uint32,    8
+        optional :if_not_modified, :bool,      9
+        optional :if_none_match,   :bool,      10
+        optional :return_head,     :bool,      11
+      end
+
+      class RpbPutResp
+        include Beefcake::Message
+        repeated :content, RpbContent, 1
+        optional :vclock,  :bytes,     2
+        optional :key,     :bytes,     3
+      end
+
+      class RpbDelReq
+        include Beefcake::Message
+        required :bucket, :bytes,  1
+        required :key,    :bytes,  2
+        optional :rw,     :uint32, 3
+        optional :vclock, :bytes,  4
+        optional :r,      :uint32, 5
+        optional :w,      :uint32, 6
+        optional :pr,     :uint32, 7
+        optional :pw,     :uint32, 8
+        optional :dw,     :uint32, 9
+      end
+
+      class RpbListBucketsResp
+        include Beefcake::Message
+        repeated :buckets, :bytes, 1
+      end
+
+      class RpbListKeysReq
+        include Beefcake::Message
+        required :bucket, :bytes, 1
+      end
+
+      class RpbListKeysResp
+        include Beefcake::Message
+        repeated :keys, :bytes, 1
+        optional :done, :bool,  2
+      end
+
+      class RpbGetBucketReq
+        include Beefcake::Message
+        required :bucket, :bytes, 1
+      end
+
+      class RpbGetBucketResp
+        include Beefcake::Message
+        required :props, RpbBucketProps, 1
+      end
+
+      class RpbSetBucketReq
+        include Beefcake::Message
+        required :bucket, :bytes,         1
+        required :props,  RpbBucketProps, 2
+      end
+
+      class RpbMapRedReq
+        include Beefcake::Message
+        required :request,      :bytes, 1
+        required :content_type, :bytes, 2
+      end
+
+      class RpbMapRedResp
+        include Beefcake::Message
+        optional :phase,    :uint32, 1
+        optional :response, :bytes,  2
+        optional :done,     :bool,   3
+      end
+
+      class RpbIndexReq
+        include Beefcake::Message
+        module IndexQueryType
+          EQ = 0
+          RANGE = 1
+        end
+
+        required :bucket,    :bytes,         1
+        required :index,     :bytes,         2
+        required :qtype,     IndexQueryType, 3
+        optional :key,       :bytes,         4
+        optional :range_min, :bytes,         5
+        optional :range_max, :bytes,         6
+      end
+
+      class RpbIndexResp
+        include Beefcake::Message
+        repeated :keys, :bytes, 1
+      end
+
+      class RpbSearchDoc
+        include Beefcake::Message
+        # We have to name this differently than the .proto file does
+        # because Beefcake uses 'fields' as an instance method.
+        repeated :properties, RpbPair, 1
+      end
+
+      class RpbSearchQueryReq
+        include Beefcake::Message
+        required :q,       :bytes,  1
+        required :index,   :bytes,  2
+        optional :rows,    :uint32, 3
+        optional :start,   :uint32, 4
+        optional :sort,    :bytes,  5
+        optional :filter,  :bytes,  6
+        optional :df,      :bytes,  7
+        optional :op,      :bytes,  8
+        repeated :fl,      :bytes,  9
+        optional :presort, :bytes, 10
+      end
+
+      class RpbSearchQueryResp
+        include Beefcake::Message
+        repeated :docs, RpbSearchDoc, 1, :default => []
+        optional :max_score, :float,  2
+        optional :num_found, :uint32, 3
+      end
+    end
+  end
+end