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

-module(riak_search_kv_hook).
-export([install/1,
         precommit_def/0,
         precommit/1]).
-ifdef(TEST).
-include_lib("eunit/include/eunit.hrl").
-export([run_mod_fun_extract_test_fun/2]).
-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 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     
install(Bucket) -> 
    BucketProps = riak_core_bucket:get_bucket(Bucket),

    %% Get the current precommit hook
    case proplists:get_value(precommit, BucketProps, []) of
        X when is_list(X) ->
            CurrentPrecommit=X;
        {struct, _}=X ->
            CurrentPrecommit=[X]
    end,

    %% Add kv/search hook - make sure there are not duplicate entries
    IndexHook = [precommit_def()],
    CleanPrecommit = CurrentPrecommit -- IndexHook,
    case CleanPrecommit ++ IndexHook 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).

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({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}).

   
-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) ->
    throw({"cannot convert to binary", Val}).

%%
%% Index the provided riak object and return ok on success.
%%
-spec index_object(riak_object(), extractdef()) -> ok.
index_object(RiakObject, Extractor) ->
    %% Set up
    {ok, RiakClient} = riak:local_client(),
    SearchClient = riak_search_client:new(RiakClient),
    Index = make_index(RiakObject),
    DocId = make_docid(RiakObject),

    %% Check the new doc is parsable and have it ready
    NewIdxDoc = make_indexed_doc(Index, DocId, RiakObject, Extractor),

    %% If all ok, remove the old entries and index the new
    remove_old_entries(RiakClient, SearchClient, Index, DocId),
    Postings = riak_indexed_doc:postings(NewIdxDoc),
    SearchClient:index_terms(Postings),

    %% Store the indexed_doc for next time
    riak_indexed_doc:put(RiakClient, NewIdxDoc),
    ok.

%% Remove any old index entries if they exist
-spec remove_old_entries(riak_client(), search_client(), index(), docid()) -> ok.
remove_old_entries(RiakClient, SearchClient, Index, DocId) ->
    case riak_indexed_doc:get(RiakClient, Index, DocId) of
        {error, notfound} ->
            ok;
        OldIdxDoc ->
            SearchClient:delete_doc_terms(OldIdxDoc)
    end.

%% Make an indexed document under Index/DocId from the RiakObject
-spec make_indexed_doc(index(), docid(), riak_object(), extractdef()) -> ok.
make_indexed_doc(Index, DocId, RiakObject, Extractor) ->
    Fields = run_extract(RiakObject, Extractor),
    IdxDoc0 = riak_indexed_doc:new(DocId, Fields, [], Index),
    {ok, IdxDoc} = riak_indexed_doc:analyze(IdxDoc0),
    IdxDoc.
                         
-spec make_index(riak_object()) -> binary().
make_index(RiakObject) ->
    riak_object:bucket(RiakObject).

-spec make_docid(riak_object()) -> binary().
make_docid(RiakObject) ->
    riak_object:key(RiakObject).
    
%% Run the extraction function against the RiakObject to get a list of
%% search fields and data
-spec run_extract(riak_object(), extractdef()) -> search_fields().
run_extract(RiakObject, {{modfun, Mod, Fun}, Arg}) ->
    Mod:Fun(RiakObject, Arg);
run_extract(RiakObject, {{qfun, Fun}, Arg}) ->
    Fun(RiakObject, Arg);<commit_msg>Taught the kv/search hook to understand extractor definitions if provided in JSON over the HTTP interface.<commit_after>%% -------------------------------------------------------------------
%%
%% Copyright (c) 2007-2010 Basho Technologies, Inc.  All Rights Reserved.
%%
%% -------------------------------------------------------------------

-module(riak_search_kv_hook).
-export([install/1,
         precommit_def/0,
         precommit/1]).
-ifdef(TEST).
-include_lib("eunit/include/eunit.hrl").
-export([run_mod_fun_extract_test_fun/2]).
-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     
install(Bucket) -> 
    BucketProps = riak_core_bucket:get_bucket(Bucket),

    %% Get the current precommit hook
    case proplists:get_value(precommit, BucketProps, []) of
        X when is_list(X) ->
            CurrentPrecommit=X;
        {struct, _}=X ->
            CurrentPrecommit=[X]
    end,

    %% Add kv/search hook - make sure there are not duplicate entries
    IndexHook = [precommit_def()],
    CleanPrecommit = CurrentPrecommit -- IndexHook,
    case CleanPrecommit ++ IndexHook 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).

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),
    Arg = proplists:get_value(<<"arg">>, Props, undefined),
    case Source of
        undefined ->
            case Name of
                undefined ->
                    throw("javascript kv/search extractor must have name or source");
                _ ->
                    {{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) ->
    throw({"cannot convert to binary", Val}).

%%
%% Index the provided riak object and return ok on success.
%%
-spec index_object(riak_object(), extractdef()) -> ok.
index_object(RiakObject, Extractor) ->
    %% Set up
    {ok, RiakClient} = riak:local_client(),
    SearchClient = riak_search_client:new(RiakClient),
    Index = make_index(RiakObject),
    DocId = make_docid(RiakObject),

    %% Check the new doc is parsable and have it ready
    NewIdxDoc = make_indexed_doc(Index, DocId, RiakObject, Extractor),

    %% If all ok, remove the old entries and index the new
    remove_old_entries(RiakClient, SearchClient, Index, DocId),
    Postings = riak_indexed_doc:postings(NewIdxDoc),
    SearchClient:index_terms(Postings),

    %% Store the indexed_doc for next time
    riak_indexed_doc:put(RiakClient, NewIdxDoc).

%% Remove any old index entries if they exist
-spec remove_old_entries(riak_client(), search_client(), index(), docid()) -> ok.
remove_old_entries(RiakClient, SearchClient, Index, DocId) ->
    case riak_indexed_doc:get(RiakClient, Index, DocId) of
        {error, notfound} ->
            ok;
        OldIdxDoc ->
            SearchClient:delete_doc_terms(OldIdxDoc)
    end.

%% Make an indexed document under Index/DocId from the RiakObject
-spec make_indexed_doc(index(), docid(), riak_object(), extractdef()) -> idxdoc().
make_indexed_doc(Index, DocId, RiakObject, Extractor) ->
    Fields = run_extract(RiakObject, Extractor),
    IdxDoc0 = riak_indexed_doc:new(DocId, Fields, [], Index),
    {ok, IdxDoc} = riak_indexed_doc:analyze(IdxDoc0),
    IdxDoc.
                         
-spec make_index(riak_object()) -> binary().
make_index(RiakObject) ->
    riak_object:bucket(RiakObject).

-spec make_docid(riak_object()) -> binary().
make_docid(RiakObject) ->
    riak_object:key(RiakObject).
    
%% Run the extraction function against the RiakObject to get a list of
%% search fields and data
-spec run_extract(riak_object(), extractdef()) -> search_fields().
run_extract(RiakObject, {{modfun, Mod, Fun}, Arg}) ->
    Mod:Fun(RiakObject, Arg);
run_extract(RiakObject, {{qfun, Fun}, Arg}) ->
    Fun(RiakObject, Arg);
{'fraction_non_alphanumeric': 0.15775220040622884, 'fraction_numerical': 0.0025850926324859974, 'mean_word_length': 3.057942057942058, 'pattern_counts': {'":': 0, '<': 33, '<?xml version=': 0, '>': 139, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '923602', 'n_tokens_mistral': 5290, 'n_tokens_neox': 5147, 'n_words': 1206}