user_cleanup.escript

#!/usr/bin/env escript
%% -*- coding: utf-8 -*-
%%! -pa /usr/lib64/riak-cs/lib/riakc-1.3.1.1/ebin /usr/lib64/riak-cs/lib/riak_pb-1.3.0/ebin /usr/lib64/riak-cs/lib/protobuffs-0.8.0/ebin /usr/lib/riak-cs/lib/riakc-1.3.1.1/ebin /usr/lib/riak-cs/lib/riak_pb-1.3.0/ebin /usr/lib/riak-cs/lib/protobuffs-0.8.0/ebin /usr/lib/riak-cs/ebin


-include_lib("riak_cs/include/riak_cs.hrl").

-define(USERS_BUCKET, <<"moss.users">>).
-define(DEFAULT_RIAK_IP, "127.0.0.1").
-define(DEFAULT_RIAK_PORT, 8087).

usage() ->
    io:format("Usage: ./user_cleanup.escript <UserId> [<RiakIP>] [<RiakPort>]~n").

main(Args) when length(Args) =:= 1 ->
    main(Args ++ [?DEFAULT_RIAK_IP, ?DEFAULT_RIAK_PORT]);
main(Args) when length(Args) < 3 ->
    usage();
main(Args) when length(Args) > 3 ->
    main(lists:sublist(Args, 3));
main(Args) when length(Args) =:= 3 ->
    [UserId, RiakHost, Port] = Args,
    RiakPort = format_port(Port),
    {ok, Pid} = riakc_pb_socket:start(RiakHost, RiakPort),
    {ok, {Record, Object}} = riak_cs_utils:get_user(UserId, Pid),
    Buckets = Record#rcs_user_v2.buckets,
    NoDeletedBuckets = [B || B <- Buckets, B#moss_bucket_v1.last_action =/= deleted],
    UpdatedUserRecord = Record#rcs_user_v2{buckets=NoDeletedBuckets},
    riak_cs_utils:save_user(UpdatedUserRecord, Object, Pid),
    riakc_pb_socket:stop(Pid),
    io:format("User record for ~p updated~n", [UserId]).

format_port(Port) when is_integer(Port) ->
    Port;
format_port(Port) when is_list(Port) ->
    list_to_integer(Port).