ListCask redux

listcask.erl

%% These are shell functions that can be used to output a keylist from a bitcask file in terms of Riak {Bucket,Key} tuples
%% They should be run from `riak attach`.  Paste in the contents and run ListAll("filename.txt").  Profit.
ListCask = fun(Vnode,OutFile) ->
        Bk_to_tuple = 
          fun(<<1:7, HasType:1, Sz:16/integer,
                     TypeOrBucket:Sz/bytes, Rest/binary>>) ->
            case HasType of
                0 ->
                    %% no type, first field is bucket
                    {TypeOrBucket, Rest};
                1 ->
                    %% has a tyoe, extract bucket as well
                    <<BucketSz:16/integer, Bucket:BucketSz/bytes, Key/binary>> = Rest,
                    {{TypeOrBucket, Bucket}, Key}
            end;
          fun(<<131:8,_Rest/binary>> = BK) ->
            binary_to_term(BK)
        end


        Fold_data = fun(Handle) ->
                FoldFun = fun(<<_:8,BucketSz:16,Bucket:BucketSz/bytes,Key/binary>>, _Ts, _PosInfo, {ok,Acc}) -> io:format(OutFile, "~5000p.~n",[{Bucket,Key}]),{ok,Acc + 1} end,
                try bitcask_fileops:fold_keys(Handle, FoldFun, {ok,0}, datafile) of
                    {ok,A} -> A;
                    {error,E} -> {error, E}
                catch
                    _Error ->
                       {error, io_lib:format("~s~s~n",["Corrupted datafile: ",bitcask_fileops:datafile_name(Handle)])}
               end
       end,

       Test_data = fun(Datafile) ->
               case catch bitcask_fileops:open_file(Datafile) of
                   {ok, Handle} ->
                       KeyCount = Fold_data(Handle),
                       bitcask_fileops:close(Handle),
                       {ok, KeyCount};
                   {error, Reason} ->
                       {error, io_lib:format("Couldn't open bitcask data file ~s with reason ~p~n", [Datafile, Reason])};
                   _Error ->
                       {error, io_lib:format("Couldn't open bitcask data file ~s~n", [Datafile])}
               end
       end,

       List_bitcask = fun(Dir) ->
               case file:list_dir(Dir) of
                   {ok, Files} ->
                       [ Test_data(filename:absname(F, Dir)) || F <- Files,
                           filename:extension(F) == ".data"];
                   {error, Reason} ->
                       {error, io_lib:format("Error listing directory ~s: ~s",[Dir, Reason])};
                   Dirlist -> {error, Dirlist}
               end
       end,


       {ok, Dpath} = application:get_env(bitcask,data_root),
       case List_bitcask(filename:absname(Vnode,Dpath)) of
         {error, Reason} ->
                io:format("Error: ~s~n",[Reason]);
         Rawlist ->
                try lists:flatten(Rawlist) of
                  A -> A
                catch
                 A -> {error, A}
                end
        end
end.

ListAll = fun(OutFileName) -> 
        {ok,Fhandle} = file:open(OutFileName,[write]),
        Rawcounts = [ ListCask( io_lib:format("~B",[P]), Fhandle ) || {riak_kv_vnode, P, _Pid} <- riak_core_vnode_manager:all_vnodes() ], 
        file:close(Fhandle),
        {KeyCount, Errors} = lists:foldl(fun(Count,{Total,Errors}) -> case Count of {ok, C} -> {Total + C, Errors}; _ -> {Total, [ Count | Errors]} end end,{0,[]}, lists:flatten(Rawcounts)),
        io:format("Listed ~B keys to file ~s~n",[KeyCount,OutFileName]),
        [ io:format("Error: ~p~n",[Err]) || Err <- Errors ],
        ok         
end.