block map-reducing over a bucket or type,bucket

Author: Luke Bakken

src/riakc_pb_socket.erl

@@ -752,6 +752,17 @@ mapred_stream(Pid, {index,Bucket,Name,StartKey,EndKey}, Query, ClientPid, Timeou
     BinEndKey = list_to_binary(integer_to_list(EndKey)),
     mapred_stream(Pid, {index,Bucket,Name,StartKey,BinEndKey}, Query, ClientPid, Timeout, CallTimeout);
 mapred_stream(Pid, Inputs, Query, ClientPid, Timeout, _CallTimeout) ->
+    AllowListing = riakc_utils:get_allow_listing(),
+    do_mapred_stream(AllowListing, Pid, Inputs, Query, ClientPid, Timeout).
+
+do_mapred_stream(false, _Pid, Bucket, _Query, _ClientPid, _Timeout) when is_binary(Bucket) ->
+    {error, <<"Bucket list operations are expensive "
+              "and should not be used in production.">>};
+do_mapred_stream(false, _Pid, {Type, Bucket}, _Query, _ClientPid, _Timeout)
+  when is_binary(Type), is_binary(Bucket) ->
+    {error, <<"Bucket list operations are expensive "
+              "and should not be used in production.">>};
+do_mapred_stream(_AllowListing, Pid, Inputs, Query, ClientPid, Timeout) ->
     MapRed = [{'inputs', Inputs},
               {'query', Query},
               {'timeout', Timeout}],

src/riakc_utils.erl

@@ -24,6 +24,7 @@
 
 -export([wait_for_list/1,
          characters_to_unicode_binary/1,
+         get_allow_listing/0,
          get_allow_listing/1]).
 
 -spec wait_for_list(non_neg_integer()) -> {ok, list()} | {error, any()}.
@@ -54,6 +55,13 @@ characters_to_unicode_binary(String) ->
 
 %% @doc Return the value of allow_listing, which, if set to 'true`
 %%      will allow listing keys and buckets.
+-spec get_allow_listing() -> boolean().
+get_allow_listing() ->
+    case application:get_env(riakc, allow_listing) of
+        {ok, true} -> true;
+        _ -> false
+    end.
+
 -spec get_allow_listing(proplists:proplist()) -> boolean().
 get_allow_listing(Options) ->
     case application:get_env(riakc, allow_listing) of

test/riakc_pb_socket_tests.erl

@@ -37,6 +37,16 @@ listing_is_blocked_test() ->
     ?assertMatch({error, E}, riakc_pb_socket:list_keys(self(), <<"b">>)),
     application:set_env(riakc, allow_listing, true).
 
+mapred_over_bucket_is_blocked_test() ->
+    application:set_env(riakc, allow_listing, false),
+    Pid = self(),
+    Input1 = <<"bucket">>,
+    Input2 = {<<"type">>, <<"bucket">>},
+    E = <<"Bucket list operations are expensive and should not be used in production.">>,
+    ?assertMatch({error, E}, riakc_pb_socket:mapred(Pid, Input1, [])),
+    ?assertMatch({error, E}, riakc_pb_socket:mapred(Pid, Input2, [])),
+    application:set_env(riakc, allow_listing, true).
+
 bad_connect_test() ->
     %% Start with an unlikely port number
     ?assertMatch({error, {tcp, econnrefused}}, riakc_pb_socket:start({127,0,0,1}, 65535)).