Skip to content

Commit 8ae580b

Browse files
lukebakkenlukebakken
authored
Merge pull request #349 from basho/features/lrb/gh-345_CLIENTS-1067
Implement errors for listing operations.
2 parents 82ae32b + 2a392af commit 8ae580b

File tree

5 files changed

+115
-18
lines changed

5 files changed

+115
-18
lines changed

‎src/riakc_pb_socket.erl‎

Lines changed: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -474,6 +474,13 @@ stream_list_buckets(Pid, Options) ->
474474
stream_list_buckets(Pid, <<"default">>, Options).
475475

476476
stream_list_buckets(Pid, Type, Options) ->
477+
AllowListing = riakc_utils:get_allow_listing(Options),
478+
do_stream_list_buckets(AllowListing, Pid, Type, Options).
479+
480+
do_stream_list_buckets(false, _Pid, _Type, _Options) ->
481+
{error, <<"Bucket and key list operations are expensive "
482+
"and should not be used in production.">>};
483+
do_stream_list_buckets(true, Pid, Type, Options) ->
477484
ST = case proplists:get_value(timeout, Options) of
478485
undefined -> ?DEFAULT_PB_TIMEOUT;
479486
T -> T
@@ -483,7 +490,14 @@ stream_list_buckets(Pid, Type, Options) ->
483490
Req = #rpblistbucketsreq{timeout=ST, type=Type, stream=true},
484491
call_infinity(Pid, {req, Req, CT, {ReqId, self()}}).
485492

486-
legacy_list_buckets(Pid, Options) ->
493+
legacy_list_buckets(Pid, Options) when is_pid(Pid), is_list(Options) ->
494+
AllowListing = riakc_utils:get_allow_listing(Options),
495+
do_legacy_list_buckets(AllowListing, Pid, Options).
496+
497+
do_legacy_list_buckets(false, _Pid, _Options) ->
498+
{error, <<"Bucket and key list operations are expensive "
499+
"and should not be used in production.">>};
500+
do_legacy_list_buckets(true, Pid, Options) ->
487501
ST = case proplists:get_value(timeout, Options) of
488502
undefined -> ?DEFAULT_PB_TIMEOUT;
489503
T -> T
@@ -542,6 +556,13 @@ stream_list_keys(Pid, Bucket, infinity) ->
542556
stream_list_keys(Pid, Bucket, Timeout) when is_integer(Timeout) ->
543557
stream_list_keys(Pid, Bucket, [{timeout, Timeout}]);
544558
stream_list_keys(Pid, Bucket, Options) ->
559+
AllowListing = riakc_utils:get_allow_listing(Options),
560+
do_stream_list_keys(AllowListing, Pid, Bucket, Options).
561+
562+
do_stream_list_keys(false, _Pid, _Bucket, _Options) ->
563+
{error, <<"Bucket and key list operations are expensive "
564+
"and should not be used in production.">>};
565+
do_stream_list_keys(true, Pid, Bucket, Options) ->
545566
ST = case proplists:get_value(timeout, Options) of
546567
undefined -> ?DEFAULT_PB_TIMEOUT;
547568
T -> T
@@ -731,6 +752,17 @@ mapred_stream(Pid, {index,Bucket,Name,StartKey,EndKey}, Query, ClientPid, Timeou
731752
BinEndKey = list_to_binary(integer_to_list(EndKey)),
732753
mapred_stream(Pid, {index,Bucket,Name,StartKey,BinEndKey}, Query, ClientPid, Timeout, CallTimeout);
733754
mapred_stream(Pid, Inputs, Query, ClientPid, Timeout, _CallTimeout) ->
755+
AllowListing = riakc_utils:get_allow_listing(),
756+
do_mapred_stream(AllowListing, Pid, Inputs, Query, ClientPid, Timeout).
757+
758+
do_mapred_stream(false, _Pid, Bucket, _Query, _ClientPid, _Timeout) when is_binary(Bucket) ->
759+
{error, <<"Bucket list operations are expensive "
760+
"and should not be used in production.">>};
761+
do_mapred_stream(false, _Pid, {Type, Bucket}, _Query, _ClientPid, _Timeout)
762+
when is_binary(Type), is_binary(Bucket) ->
763+
{error, <<"Bucket list operations are expensive "
764+
"and should not be used in production.">>};
765+
do_mapred_stream(_AllowListing, Pid, Inputs, Query, ClientPid, Timeout) ->
734766
MapRed = [{'inputs', Inputs},
735767
{'query', Query},
736768
{'timeout', Timeout}],

‎src/riakc_ts.erl‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -224,6 +224,13 @@ stream_list_keys(Pid, Table, Timeout) when is_integer(Timeout) ->
224224
stream_list_keys(Pid, Table, [{timeout, Timeout}]);
225225
stream_list_keys(Pid, Table, Options)
226226
when is_pid(Pid), (is_binary(Table) orelse is_list(Table)), is_list(Options) ->
227+
AllowListing = riakc_utils:get_allow_listing(Options),
228+
do_stream_list_keys(AllowListing, Pid, Table, Options).
229+
230+
do_stream_list_keys(false, _Pid, _Table, _Options) ->
231+
{error, <<"Bucket and key list operations are expensive "
232+
"and should not be used in production.">>};
233+
do_stream_list_keys(true, Pid, Table, Options) ->
227234
T = riakc_utils:characters_to_unicode_binary(Table),
228235
ReqTimeout = proplists:get_value(timeout, Options),
229236
Req = #tslistkeysreq{table = T, timeout = ReqTimeout},

‎src/riakc_utils.erl‎

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,10 @@
2222

2323
-module(riakc_utils).
2424

25-
-export([wait_for_list/1, characters_to_unicode_binary/1]).
25+
-export([wait_for_list/1,
26+
characters_to_unicode_binary/1,
27+
get_allow_listing/0,
28+
get_allow_listing/1]).
2629

2730
-spec wait_for_list(non_neg_integer()) -> {ok, list()} | {error, any()}.
2831
%% @doc Wait for the results of a listing operation
@@ -49,3 +52,23 @@ characters_to_unicode_binary(String) ->
4952
Binary ->
5053
Binary
5154
end.
55+
56+
%% @doc Return the value of allow_listing, which, if set to 'true`
57+
%% will allow listing keys and buckets.
58+
-spec get_allow_listing() -> boolean().
59+
get_allow_listing() ->
60+
case application:get_env(riakc, allow_listing) of
61+
{ok, true} -> true;
62+
_ -> false
63+
end.
64+
65+
-spec get_allow_listing(proplists:proplist()) -> boolean().
66+
get_allow_listing(Options) ->
67+
case application:get_env(riakc, allow_listing) of
68+
{ok, true} -> true;
69+
_ ->
70+
case proplists:get_value(allow_listing, Options) of
71+
true -> true;
72+
_ -> false
73+
end
74+
end.

‎test/riakc_pb_socket_tests.erl‎

Lines changed: 39 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -30,24 +30,45 @@
3030
-include_lib("riak_pb/include/riak_pb_kv_codec.hrl").
3131
-include_lib("eunit/include/eunit.hrl").
3232

33+
listing_is_blocked_test() ->
34+
application:set_env(riakc, allow_listing, false),
35+
E = <<"Bucket and key list operations are expensive and should not be used in production.">>,
36+
?assertMatch({error, E}, riakc_pb_socket:list_buckets(self())),
37+
?assertMatch({error, E}, riakc_pb_socket:list_keys(self(), <<"b">>)),
38+
application:set_env(riakc, allow_listing, true).
39+
40+
mapred_over_bucket_is_blocked_test() ->
41+
application:set_env(riakc, allow_listing, false),
42+
Pid = self(),
43+
Input1 = <<"bucket">>,
44+
Input2 = {<<"type">>, <<"bucket">>},
45+
E = <<"Bucket list operations are expensive and should not be used in production.">>,
46+
?assertMatch({error, E}, riakc_pb_socket:mapred(Pid, Input1, [])),
47+
?assertMatch({error, E}, riakc_pb_socket:mapred(Pid, Input2, [])),
48+
application:set_env(riakc, allow_listing, true).
49+
3350
bad_connect_test() ->
3451
%% Start with an unlikely port number
35-
?assertEqual({error, {tcp, econnrefused}}, riakc_pb_socket:start({127,0,0,1}, 65535)).
52+
?assertMatch({error, {tcp, econnrefused}}, riakc_pb_socket:start({127,0,0,1}, 65535)).
3653

3754
queue_disconnected_test() ->
55+
application:set_env(riakc, allow_listing, true),
3856
%% Start with an unlikely port number
3957
{ok, Pid} = riakc_pb_socket:start({127,0,0,1}, 65535, [queue_if_disconnected]),
40-
?assertEqual({error, timeout}, riakc_pb_socket:ping(Pid, 10)),
41-
?assertEqual({error, timeout}, riakc_pb_socket:list_keys(Pid, <<"b">>, 10)),
42-
riakc_pb_socket:stop(Pid).
58+
?assertMatch({error, timeout}, riakc_pb_socket:ping(Pid, 10)),
59+
?assertMatch({error, timeout}, riakc_pb_socket:list_keys(Pid, <<"b">>, 10)),
60+
riakc_pb_socket:stop(Pid),
61+
application:set_env(riakc, allow_listing, false).
4362

4463
auto_reconnect_bad_connect_test() ->
64+
application:set_env(riakc, allow_listing, true),
4565
%% Start with an unlikely port number
4666
{ok, Pid} = riakc_pb_socket:start({127,0,0,1}, 65535, [auto_reconnect]),
47-
?assertEqual({false, []}, riakc_pb_socket:is_connected(Pid)),
48-
?assertEqual({error, disconnected}, riakc_pb_socket:ping(Pid)),
49-
?assertEqual({error, disconnected}, riakc_pb_socket:list_keys(Pid, <<"b">>)),
50-
riakc_pb_socket:stop(Pid).
67+
?assertMatch({false, []}, riakc_pb_socket:is_connected(Pid)),
68+
?assertMatch({error, disconnected}, riakc_pb_socket:ping(Pid)),
69+
?assertMatch({error, disconnected}, riakc_pb_socket:list_keys(Pid, <<"b">>)),
70+
riakc_pb_socket:stop(Pid),
71+
application:set_env(riakc, allow_listing, false).
5172

5273
server_closes_socket_test() ->
5374
%% Silence SASL junk when socket closes.
@@ -69,7 +90,7 @@ server_closes_socket_test() ->
6990
ok = gen_tcp:close(Listen),
7091
receive
7192
Msg1 -> % result of ping from spawned process above
72-
?assertEqual({error, disconnected}, Msg1)
93+
?assertMatch({error, disconnected}, Msg1)
7394
end,
7495
%% Wait for spawned process to exit
7596
Mref = erlang:monitor(process, Pid),
@@ -96,7 +117,7 @@ auto_reconnect_server_closes_socket_test() ->
96117
ok = gen_tcp:close(Listen),
97118
receive
98119
Msg ->
99-
?assertEqual({error, disconnected}, Msg)
120+
?assertMatch({error, disconnected}, Msg)
100121
end,
101122
%% Server will not have had a chance to reconnect yet, reason counters empty.
102123
?assertMatch({false, []}, riakc_pb_socket:is_connected(Pid)),
@@ -139,8 +160,8 @@ integration_tests() ->
139160
[{"ping",
140161
?_test( begin
141162
{ok, Pid} = riakc_test_utils:start_link(),
142-
?assertEqual(pong, riakc_pb_socket:ping(Pid)),
143-
?assertEqual(true, riakc_pb_socket:is_connected(Pid)),
163+
?assertMatch(pong, riakc_pb_socket:ping(Pid)),
164+
?assertMatch(true, riakc_pb_socket:is_connected(Pid)),
144165
riakc_pb_socket:stop(Pid)
145166
end)},
146167

@@ -1453,9 +1474,13 @@ integration_test_() ->
14531474
SetupFun = fun() ->
14541475
%% Grab the riakclient_pb.proto file
14551476
code:add_pathz("../ebin"),
1456-
ok = riakc_test_utils:maybe_start_network()
1477+
ok = riakc_test_utils:maybe_start_network(),
1478+
application:set_env(riakc, allow_listing, true)
14571479
end,
1458-
CleanupFun = fun(_) -> net_kernel:stop() end,
1480+
CleanupFun = fun(_) ->
1481+
net_kernel:stop(),
1482+
application:set_env(riakc, allow_listing, false)
1483+
end,
14591484
GenFun = fun() ->
14601485
case catch net_adm:ping(riakc_test_utils:test_riak_node()) of
14611486
pong -> integration_tests();

‎test/riakc_ts_tests.erl‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,12 @@
3636
-define(FIVE_MINS_AGO, 1443796600987).
3737
-define(NOW, 1443796900987).
3838

39+
listing_is_blocked_test() ->
40+
application:set_env(riakc, allow_listing, false),
41+
E = <<"Bucket and key list operations are expensive and should not be used in production.">>,
42+
?assertMatch({error, E}, riakc_ts:stream_list_keys(self(), <<"b">>)),
43+
application:set_env(riakc, allow_listing, true).
44+
3945
integration_tests({ok, _Props}) ->
4046
[{"ping",
4147
?_test(begin
@@ -86,9 +92,13 @@ integration_test_() ->
8692
SetupFun = fun() ->
8793
%% Grab the riakclient_pb.proto file
8894
code:add_pathz("../ebin"),
89-
ok = riakc_test_utils:maybe_start_network()
95+
ok = riakc_test_utils:maybe_start_network(),
96+
application:set_env(riakc, allow_listing, true)
9097
end,
91-
CleanupFun = fun(_) -> net_kernel:stop() end,
98+
CleanupFun = fun(_) ->
99+
net_kernel:stop(),
100+
application:set_env(riakc, allow_listing, false)
101+
end,
92102
GenFun = fun() ->
93103
case catch net_adm:ping(riakc_test_utils:test_riak_node()) of
94104
pong -> generate_integration_tests();

0 commit comments

Comments
 (0)