<commit_before>%% -------------------------------------------------------------------
%%
%% Copyright (c) 2007-2010 Basho Technologies, Inc.  All Rights Reserved.
%%
%% -------------------------------------------------------------------

-module(riak_search_kv_hook).
-export([install/1,
         uninstall/1,
         precommit_def/0,
         precommit/1]).

-include("riak_search.hrl").

-ifdef(TEST).
-include_lib("eunit/include/eunit.hrl").
-export([run_mod_fun_extract_test_fun/3]).
-endif.

-define(DEFAULT_EXTRACTOR, {modfun, riak_search_kv_extractor, extract}).
-define(DEFAULT_ARGS,      undefined).

-type user_funterm() :: {modfun, user_modname(), user_funname()} |
                        {qfun, extract_qfun()} |
                        {jsanon, user_strorbin()} |
                        {jsanon, {user_strorbin(), user_strorbin()}} |
                        {jsfun, user_strorbin()}.
-type user_modname() :: string() | module().
-type user_funname() :: string() | atom().
-type user_strorbin() :: string() | binary().                       


-type extractdef() :: {funterm(), args()}.
-type funterm() :: {modfun, atom(), atom()} |
                   {qfun, extract_qfun()} |
                   {jsanon, binary()} |
                   {jsanon, {binary(), binary()}} |
                   {jsfun, binary()}.

-type riak_client() :: tuple(). % no good way to define riak_client
-type riak_object() :: tuple(). % no good way to define riak_object
-type search_client() :: tuple().       

-type extract_qfun() :: fun((riak_object(),any()) -> search_fields()).
-type args() :: any().

-type index() :: binary().
-type docid() :: binary().
-type idxdoc() :: tuple(). % #riak_indexed_doc{}

-type search_fields() :: [{search_field(),search_data()}].
-type search_field() :: string().
-type search_data() :: string() | binary().
    

%% Install the kv/search integration hook on the specified bucket
%% TODO: The code below can be
%% simplified. riak_core_bucket:set_bucket/2 does not require you to
%% set ALL bucket props, just the properties that have changed.
install(Bucket) -> 
    BucketProps = riak_core_bucket:get_bucket(Bucket),
    CleanPrecommit = strip_precommit(BucketProps),
    case CleanPrecommit ++ [precommit_def()] of
        [{struct, _}]=Y ->
            UpdPrecommit=Y;
        Y ->
            UpdPrecommit=Y
    end,

    %% Update the bucket properties
    UpdBucketProps = lists:keyreplace(precommit, 1, BucketProps, 
                                      {precommit, UpdPrecommit}),
    riak_core_bucket:set_bucket(Bucket, UpdBucketProps).

%% Uninstall kv/search integration hook from specified bucket
uninstall(Bucket) ->
    BucketProps = riak_core_bucket:get_bucket(Bucket),
    CleanPrecommit = strip_precommit(BucketProps),

    %% Update the bucket properties
    UpdBucketProps = lists:keyreplace(precommit, 1, BucketProps, 
                                      {precommit, CleanPrecommit}),
    riak_core_bucket:set_bucket(Bucket, UpdBucketProps).

precommit_def() ->
    {struct, [{<<"mod">>,atom_to_binary(?MODULE, latin1)},
              {<<"fun">>,<<"precommit">>}]}.


%% Precommit hook for riak k/v and search integration.  Executes
%% the desired mapping on the riak object to produce a search
%% document to store in riak search.
%%
-spec precommit(riak_object()) -> {fail, any()} | riak_object().
precommit(RiakObject) ->
    Extractor = get_extractor(RiakObject),
    try
        case index_object(RiakObject, Extractor) of
            ok ->
                RiakObject;
            {error, Reason1} ->
                {fail, Reason1}
        end
    catch
        throw:Reason2 ->
            {fail, Reason2}
    end.

%% Decide if an object should be indexed, and if so the extraction function to 
%% pull out the search fields.
-spec get_extractor(riak_object()) -> {funterm(), any()}.
get_extractor(RiakObject) ->
    BucketProps = riak_core_bucket:get_bucket(riak_object:bucket(RiakObject)),
    validate_extractor(proplists:get_value(rs_extractfun, BucketProps, undefined)).

%% Validate the extraction function and normalize to {FunTerm, Args}
-spec validate_extractor(undefined |
                         user_funterm() |
                         {user_funterm(), args()}) -> {funterm(), args()}.
validate_extractor(undefined) ->
    {?DEFAULT_EXTRACTOR, ?DEFAULT_ARGS};
validate_extractor({struct, JsonExtractor}) ->
    Lang = proplists:get_value(<<"language">>, JsonExtractor),    
    validate_extractor(erlify_json_funterm(Lang, JsonExtractor));
validate_extractor({FunTerm, Args}) when is_tuple(FunTerm) ->
    {validate_funterm(FunTerm), Args};
validate_extractor(FunTerm) ->
    {validate_funterm(FunTerm), undefined}.

-spec validate_funterm(user_funterm()) -> funterm().
validate_funterm({modfun, Mod, Fun}) ->
    {modfun, to_modfun(Mod), to_modfun(Fun)};
validate_funterm({qfun, Fun}=FunTerm) when is_function(Fun) ->
    FunTerm;
validate_funterm({jsanon, {Bucket, Key}}) ->
    {jsanon, {to_binary(Bucket), to_binary(Key)}};
validate_funterm({jsanon, Source}) ->
    {jsanon, to_binary(Source)};
validate_funterm({jsfun, Name}) ->
    {jsfun, to_binary(Name)};
validate_funterm(FunTerm) ->
    throw({"cannot parse funterm", FunTerm}).

%% Decode a bucket property that was set using JSON/HTTP interface
erlify_json_funterm(<<"erlang">>, Props) ->
    Mod = to_modfun(proplists:get_value(<<"module">>, Props, undefined)),
    Fun = to_modfun(proplists:get_value(<<"function">>, Props, undefined)),
    Arg = proplists:get_value(<<"arg">>, Props, undefined),
    {{modfun, Mod, Fun}, Arg};
erlify_json_funterm(<<"javascript">>, Props) ->
    Source = proplists:get_value(<<"source">>, Props, undefined),
    Name = proplists:get_value(<<"name">>, Props, undefined),
    Bucket = proplists:get_value(<<"bucket">>, Props, undefined),
    Key = proplists:get_value(<<"key">>, Props, undefined),
    Arg = proplists:get_value(<<"arg">>, Props, undefined),
    case Source of
        undefined ->
            case Name of
                undefined ->
                    case (Bucket == undefined) orelse (Key == undefined) of
                        true ->
                            throw("javascript kv/search extractor must have"
                                  "name, source, or bucket and key");
                        _ ->
                            {{jsanon, {Bucket, Key}}, Arg}
                    end;
                _ ->
                    {{jsfun, Name}, Arg}
            end;
        _ ->
            {{jsanon, Source}, Arg}
    end;
erlify_json_funterm(Lang, _Props) ->
    throw({"kv/search extractors must be written in erlang or javascript", Lang}).

     

-spec to_modfun(list() | atom()) -> atom().
to_modfun(List) when is_list(List) ->
    %% Using list_to_atom here so that the extractor module
    %% does not need to be pre-loaded.  
    list_to_atom(List);
to_modfun(Atom) when is_atom(Atom) ->
    Atom;
to_modfun(Val) ->
    throw({"cannot convert to module/function name", Val}).
   
-spec to_binary(atom() | string() | binary()) -> binary().
to_binary(Atom) when is_atom(Atom) ->
    atom_to_binary(Atom, utf8);
to_binary(List) when is_list(List) ->
    list_to_binary(List);
to_binary(Bin) when is_binary(Bin) ->
    Bin;
to_binary(Val) -><commit_msg>Fix riak_search_kv_hook:to_modfun/1 to handle binaries.<commit_after>%% -------------------------------------------------------------------
%%
%% Copyright (c) 2007-2010 Basho Technologies, Inc.  All Rights Reserved.
%%
%% -------------------------------------------------------------------

-module(riak_search_kv_hook).
-export([install/1,
         uninstall/1,
         precommit_def/0,
         precommit/1]).

-include("riak_search.hrl").

-ifdef(TEST).
-include_lib("eunit/include/eunit.hrl").
-export([run_mod_fun_extract_test_fun/3]).
-endif.

-define(DEFAULT_EXTRACTOR, {modfun, riak_search_kv_extractor, extract}).
-define(DEFAULT_ARGS,      undefined).

-type user_funterm() :: {modfun, user_modname(), user_funname()} |
                        {qfun, extract_qfun()} |
                        {jsanon, user_strorbin()} |
                        {jsanon, {user_strorbin(), user_strorbin()}} |
                        {jsfun, user_strorbin()}.
-type user_modname() :: string() | module().
-type user_funname() :: string() | atom().
-type user_strorbin() :: string() | binary().                       


-type extractdef() :: {funterm(), args()}.
-type funterm() :: {modfun, atom(), atom()} |
                   {qfun, extract_qfun()} |
                   {jsanon, binary()} |
                   {jsanon, {binary(), binary()}} |
                   {jsfun, binary()}.

-type riak_client() :: tuple(). % no good way to define riak_client
-type riak_object() :: tuple(). % no good way to define riak_object
-type search_client() :: tuple().       

-type extract_qfun() :: fun((riak_object(),any()) -> search_fields()).
-type args() :: any().

-type index() :: binary().
-type docid() :: binary().
-type idxdoc() :: tuple(). % #riak_indexed_doc{}

-type search_fields() :: [{search_field(),search_data()}].
-type search_field() :: string().
-type search_data() :: string() | binary().
    

%% Install the kv/search integration hook on the specified bucket
%% TODO: The code below can be
%% simplified. riak_core_bucket:set_bucket/2 does not require you to
%% set ALL bucket props, just the properties that have changed.
install(Bucket) -> 
    BucketProps = riak_core_bucket:get_bucket(Bucket),
    CleanPrecommit = strip_precommit(BucketProps),
    case CleanPrecommit ++ [precommit_def()] of
        [{struct, _}]=Y ->
            UpdPrecommit=Y;
        Y ->
            UpdPrecommit=Y
    end,

    %% Update the bucket properties
    UpdBucketProps = lists:keyreplace(precommit, 1, BucketProps, 
                                      {precommit, UpdPrecommit}),
    riak_core_bucket:set_bucket(Bucket, UpdBucketProps).

%% Uninstall kv/search integration hook from specified bucket
uninstall(Bucket) ->
    BucketProps = riak_core_bucket:get_bucket(Bucket),
    CleanPrecommit = strip_precommit(BucketProps),

    %% Update the bucket properties
    UpdBucketProps = lists:keyreplace(precommit, 1, BucketProps, 
                                      {precommit, CleanPrecommit}),
    riak_core_bucket:set_bucket(Bucket, UpdBucketProps).

precommit_def() ->
    {struct, [{<<"mod">>,atom_to_binary(?MODULE, latin1)},
              {<<"fun">>,<<"precommit">>}]}.


%% Precommit hook for riak k/v and search integration.  Executes
%% the desired mapping on the riak object to produce a search
%% document to store in riak search.
%%
-spec precommit(riak_object()) -> {fail, any()} | riak_object().
precommit(RiakObject) ->
    Extractor = get_extractor(RiakObject),
    ?PRINT(Extractor),
    try
        case index_object(RiakObject, Extractor) of
            ok ->
                RiakObject;
            {error, Reason1} ->
                {fail, Reason1}
        end
    catch
        throw:Reason2 ->
            {fail, Reason2}
    end.

%% Decide if an object should be indexed, and if so the extraction function to 
%% pull out the search fields.
-spec get_extractor(riak_object()) -> {funterm(), any()}.
get_extractor(RiakObject) ->
    BucketProps = riak_core_bucket:get_bucket(riak_object:bucket(RiakObject)),
    validate_extractor(proplists:get_value(rs_extractfun, BucketProps, undefined)).

%% Validate the extraction function and normalize to {FunTerm, Args}
-spec validate_extractor(undefined |
                         user_funterm() |
                         {user_funterm(), args()}) -> {funterm(), args()}.
validate_extractor(undefined) ->
    {?DEFAULT_EXTRACTOR, ?DEFAULT_ARGS};
validate_extractor({struct, JsonExtractor}) ->
    Lang = proplists:get_value(<<"language">>, JsonExtractor),    
    validate_extractor(erlify_json_funterm(Lang, JsonExtractor));
validate_extractor({FunTerm, Args}) when is_tuple(FunTerm) ->
    {validate_funterm(FunTerm), Args};
validate_extractor(FunTerm) ->
    {validate_funterm(FunTerm), undefined}.

-spec validate_funterm(user_funterm()) -> funterm().
validate_funterm({modfun, Mod, Fun}) ->
    {modfun, to_modfun(Mod), to_modfun(Fun)};
validate_funterm({qfun, Fun}=FunTerm) when is_function(Fun) ->
    FunTerm;
validate_funterm({jsanon, {Bucket, Key}}) ->
    {jsanon, {to_binary(Bucket), to_binary(Key)}};
validate_funterm({jsanon, Source}) ->
    {jsanon, to_binary(Source)};
validate_funterm({jsfun, Name}) ->
    {jsfun, to_binary(Name)};
validate_funterm(FunTerm) ->
    throw({"cannot parse funterm", FunTerm}).

%% Decode a bucket property that was set using JSON/HTTP interface
erlify_json_funterm(<<"erlang">>, Props) ->
    Mod = to_modfun(proplists:get_value(<<"module">>, Props, undefined)),
    Fun = to_modfun(proplists:get_value(<<"function">>, Props, undefined)),
    Arg = proplists:get_value(<<"arg">>, Props, undefined),
    {{modfun, Mod, Fun}, Arg};
erlify_json_funterm(<<"javascript">>, Props) ->
    Source = proplists:get_value(<<"source">>, Props, undefined),
    Name = proplists:get_value(<<"name">>, Props, undefined),
    Bucket = proplists:get_value(<<"bucket">>, Props, undefined),
    Key = proplists:get_value(<<"key">>, Props, undefined),
    Arg = proplists:get_value(<<"arg">>, Props, undefined),
    case Source of
        undefined ->
            case Name of
                undefined ->
                    case (Bucket == undefined) orelse (Key == undefined) of
                        true ->
                            throw("javascript kv/search extractor must have"
                                  "name, source, or bucket and key");
                        _ ->
                            {{jsanon, {Bucket, Key}}, Arg}
                    end;
                _ ->
                    {{jsfun, Name}, Arg}
            end;
        _ ->
            {{jsanon, Source}, Arg}
    end;
erlify_json_funterm(Lang, _Props) ->
    throw({"kv/search extractors must be written in erlang or javascript", Lang}).

     

-spec to_modfun(list() | atom()) -> atom().
to_modfun(List) when is_list(List) ->
    %% Using list_to_atom here so that the extractor module
    %% does not need to be pre-loaded.  
    list_to_atom(List);
to_modfun(Binary) when is_binary(Binary) ->
    to_modfun(binary_to_list(Binary));
to_modfun(Atom) when is_atom(Atom) ->
    Atom;
to_modfun(Val) ->
    throw({"cannot convert to module/function name", Val}).
   
-spec to_binary(atom() | string() | binary()) -> binary().
to_binary(Atom) when is_atom(Atom) ->
    atom_to_binary(Atom, utf8);
to_binary(List) when is_list(List) ->
    list_to_binary(List);
to_binary(Bin) when is_binary(Bin) ->
    Bin;
to_binary(Val) ->
{'fraction_non_alphanumeric': 0.1627589967284624, 'fraction_numerical': 0.0030670665212649945, 'mean_word_length': 2.741203467618562, 'pattern_counts': {'":': 0, '<': 59, '<?xml version=': 0, '>': 146, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4929282', 'n_tokens_mistral': 4607, 'n_tokens_neox': 4493, 'n_words': 1025}