Anonymous avatar Anonymous committed 92f7309

make it easier to specify the client id when creating a new riak_client, and expose the ability all the way throught he HTTP layers (X-Riak-ClientId header)

Comments (0)

Files changed (5)

 %%      links for following.
 -module(jiak).
 
--export([local_client/0, client_connect/1]).
+-export([local_client/0, local_client/1,
+         client_connect/1, client_connect/2]).
 -export([default_jiak_bucket_props/0]).
 -export([standard_sibling_merge/1]).
 
 -include_lib("eunit/include/eunit.hrl").
 
 %% @spec local_client() -> {ok, jiak_client()}|error_term()
+%% @equiv local_client(undefined)
+local_client() -> local_client(undefined).
+
+%% @spec local_client(binary()|undefined) ->
+%%         {ok, jiak_client()}|error_term()
 %% @doc Open a Riak client for modifying Jiak objects.
-%% @see riak:local_client/0
-local_client() -> client_connect(node()).
+%% @see riak:local_client/1
+local_client(ClientId) -> client_connect(node(), ClientId).
 
 %% @spec client_connect(Node :: node())
 %%        -> {ok, Client :: jiak_client()} | exception
+%% @equiv client_connect(Node, undefined)
+client_connect(Node) -> client_connect(Node, undefined).
+
+%% @spec client_connect(Node :: node(), ClientId :: binary()|undefined)
+%%        -> {ok, Client :: jiak_client()} | exception
 %% @doc The usual way to get a client.  Timeout often means either a bad
 %%      cookie or a poorly-connected distributed erlang network.
-client_connect(Node) -> 
-    case riak:client_connect(Node) of
+client_connect(Node, ClientId) ->
+    case riak:client_connect(Node, ClientId) of
         {ok, C} -> {ok, jiak_client:new(C)};
         Error -> Error
     end.

src/jiak_resource.erl

 %%      attempt to open a client to Riak, and will fail if it is
 %%      unable to do so.
 init(Props) ->
-    {ok, JiakClient} = 
+    ClientType = 
         case proplists:get_value(riak_local, Props) of
-            true ->
-                jiak:local_client();
+            true -> local;
             _ ->
                 Node = proplists:get_value(riak_node, Props),
                 Cookie = proplists:get_value(riak_cookie, Props),
                 erlang:set_cookie(node(), Cookie),
-                jiak:client_connect(Node)
+                {remote, Node}
         end,
     {ok, #ctx{jiak_name=proplists:get_value(jiak_name, Props),
               key=proplists:get_value(key_type, Props),
-              jiak_client=JiakClient}}.
+              jiak_client=ClientType}}.
 
 %% @spec service_available(webmachine:wrq(), context()) -> 
 %%           {boolean, webmachine:wrq(), context()}
 %%      same name as the bucket.  If no module is found, the bucket 
 %%      configuration metadata in the ring is used, and must contain
 %%      a valid Jiak schema.  
-service_available(ReqData, Context=#ctx{key=container}) ->
+service_available(ReqData, Context=#ctx{jiak_client=ClientType}) ->
+    {ok, Client} = case ClientType of
+                       local ->
+                           jiak:local_client(get_client_id(ReqData));
+                       {remote, Node} ->
+                           jiak:client_connect(Node, get_client_id(ReqData))
+                   end,
+    service_available2(ReqData, Context#ctx{jiak_client=Client}).
+
+%% @spec service_available2(webmachine:wrq(), context()) -> 
+%%           {boolean, webmachine:wrq(), context()}
+%% @doc Continue service_available processing after the riak_client
+%%      is created.
+service_available2(ReqData, Context=#ctx{key=container}) ->
     {ServiceAvailable, NewCtx} = 
         case wrq:method(ReqData) of
             'PUT' -> 
                 {true, Context#ctx{module=Mod}}
         end,
     {ServiceAvailable, ReqData, NewCtx};
-service_available(ReqData, Context) ->
+service_available2(ReqData, Context) ->
     Mod = jiak_util:get_jiak_module(ReqData),
     {true, ReqData, Context#ctx{module=Mod}}.
 
+%% @spec get_client_id(reqdata()) -> term()
+%% @doc Extract the request's preferred client id from the
+%%      X-Riak-ClientId header.  Return value will be:
+%%        'undefined' if no header was found
+%%        32-bit binary() if the header could be base64-decoded
+%%           into a 32-bit binary
+%%        string() if the header could not be base64-decoded
+%%           into a 32-bit binary
+get_client_id(RD) ->
+    case wrq:get_req_header("X-Riak-ClientId", RD) of
+        undefined -> undefined;
+        RawId ->
+            case catch base64:decode(RawId) of
+                ClientId= <<_:32>> -> ClientId;
+                _ -> RawId
+            end
+    end.
+
 %% @spec allowed_methods(webmachine:wrq(), context()) ->
 %%          {[http_method()], webmachine:wrq(), context()}
 %% @type http_method() = 'HEAD'|'GET'|'POST'|'PUT'|'DELETE'
 -define(HEAD_VCLOCK,   "X-Riak-Vclock").
 -define(HEAD_LINK,     "Link").
 -define(HEAD_ENCODING, "Content-Encoding").
+-define(HEAD_CLIENT,   "X-Riak-ClientId").
 
 %% Names of JSON fields in bucket properties
 -define(JSON_PROPS,   <<"props">>).

src/raw_http_resource.erl

 %%      bindings from the dispatch, as well as any vtag
 %%      query parameter.
 service_available(RD, Ctx=#ctx{riak=RiakProps}) ->
-    case get_riak_client(RiakProps) of
+    case get_riak_client(RiakProps, get_client_id(RD)) of
         {ok, C} ->
             {true,
              RD,
              Ctx}
     end.
 
-%% @spec get_riak_client(local|{node(),Cookie::atom()}) ->
+%% @spec get_riak_client(local|{node(),Cookie::atom()}, term()) ->
 %%          {ok, riak_client()} | error()
 %% @doc Get a riak_client.
-get_riak_client(local) ->
-    riak:local_client();
-get_riak_client({Node, Cookie}) ->
+get_riak_client(local, ClientId) ->
+    riak:local_client(ClientId);
+get_riak_client({Node, Cookie}, ClientId) ->
     erlang:set_cookie(node(), Cookie),
-    riak:client_connect(Node).
+    riak:client_connect(Node, ClientId).
+
+%% @spec get_client_id(reqdata()) -> term()
+%% @doc Extract the request's preferred client id from the
+%%      X-Riak-ClientId header.  Return value will be:
+%%        'undefined' if no header was found
+%%        32-bit binary() if the header could be base64-decoded
+%%           into a 32-bit binary
+%%        string() if the header could not be base64-decoded
+%%           into a 32-bit binary
+get_client_id(RD) ->
+    case wrq:get_req_header(?HEAD_CLIENT, RD) of
+        undefined -> undefined;
+        RawId ->
+            case catch base64:decode(RawId) of
+                ClientId= <<_:32>> -> ClientId;
+                _ -> RawId
+            end
+    end.
 
 %% @spec allowed_methods(reqdata(), context()) ->
 %%          {[method()], reqdata(), context()}
 -author('Bryan Fink <bryan@basho.com>').
 -export([start/0, start/1, stop/0, stop/1]).
 -export([get_app_env/1,get_app_env/2]).
--export([client_connect/1,local_client/0]).
+-export([client_connect/1,client_connect/2,
+         local_client/0,local_client/1]).
 
 -include_lib("eunit/include/eunit.hrl").
 
     end.
 
 %% @spec local_client() -> {ok, Client :: riak_client()}
+%% @equiv local_client(undefined)
+local_client() ->
+    local_client(undefined).
+
+%% @spec local_client(binary()|undefined) -> {ok, Client :: riak_client()}
 %% @doc When you want a client for use on a running Riak node.
-local_client() -> client_connect(node()).
-
+%%      ClientId should be a 32-bit binary.  If it is not, a
+%%      32-bit binary will be created from ClientId by phash2/1.
+%%      If ClientId is the atom 'undefined', a random ClientId will
+%%      be chosen.
+local_client(ClientId) ->
+    client_connect(node(), ClientId).
 
 %% @spec client_connect(Node :: node())
 %%        -> {ok, Client :: riak_client()} | {error, timeout}
+%% @equiv client_connect(Node, undefined)
+client_connect(Node) -> 
+    client_connect(Node, undefined).
+
+%% @spec client_connect(node(), binary()|undefined)
+%%         -> {ok, Client :: riak_client} | {error, timeout}
 %% @doc The usual way to get a client.  Timeout often means either a bad
 %%      cookie or a poorly-connected distributed erlang network.
-client_connect(Node) -> 
+%%      ClientId should be a 32-bit binary.  If it is not, a
+%%      32-bit binary will be created from ClientId by phash2/1.
+%%      If ClientId is the atom 'undefined', a random ClientId will
+%%      be chosen.
+client_connect(Node, ClientId= <<_:32>>) ->
     % Make sure we can reach this node...
     case net_adm:ping(Node) of
         pang -> {error, {could_not_reach_node, Node}};
-        pong -> {ok, riak_client:new(Node, riak_util:mkclientid(Node))}
-    end.
-
+        pong -> {ok, riak_client:new(Node, ClientId)}
+    end;
+client_connect(Node, undefined) ->
+    client_connect(Node, riak_util:mkclientid(Node));
+client_connect(Node, Other) ->
+    client_connect(Node, <<(erlang:phash2(Other)):32>>).
 
 %% @spec ensure_started(Application :: atom()) -> ok
 %% @doc Start the named application if not already started.
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.