Support partitioned queries in Mango

Co-authored-by: Garren Smith <garren.smith@gmail.com>
Co-authored-by: Robert Newson <rnewson@apache.org>

Author: 3 people

src/mango/src/mango_cursor.erl

@@ -71,6 +71,7 @@ explain(#cursor{}=Cursor) ->
     {[
         {dbname, mango_idx:dbname(Idx)},
         {index, mango_idx:to_json(Idx)},
+        {partitioned, mango_idx:partitioned(Idx)},
         {selector, Selector},
         {opts, {Opts}},
         {limit, Limit},

src/mango/src/mango_cursor_text.erl

@@ -77,6 +77,7 @@ explain(Cursor) ->
     } = Cursor,
     [
         {'query', mango_selector_text:convert(Selector)},
+        {partition, get_partition(Opts, null)},
         {sort, sort_query(Opts, Selector)}
     ].
 
@@ -93,6 +94,7 @@ execute(Cursor, UserFun, UserAcc) ->
     } = Cursor,
     QueryArgs = #index_query_args{
         q = mango_selector_text:convert(Selector),
+        partition = get_partition(Opts, nil),
         sort = sort_query(Opts, Selector),
         raw_bookmark = true
     },
@@ -237,6 +239,13 @@ sort_query(Opts, Selector) ->
     end.
 
 
+get_partition(Opts, Default) ->
+    case couch_util:get_value(partition, Opts) of
+        <<>> -> Default;
+        Else -> Else
+    end.
+
+
 get_bookmark(Opts) ->
     case lists:keyfind(bookmark, 1, Opts) of
         {_, BM} when is_list(BM), BM /= [] ->

src/mango/src/mango_cursor_view.erl

@@ -71,6 +71,7 @@ explain(Cursor) ->
         {include_docs, Args#mrargs.include_docs},
         {view_type, Args#mrargs.view_type},
         {reduce, Args#mrargs.reduce},
+        {partition, couch_mrview_util:get_extra(Args, partition, null)},
         {start_key, maybe_replace_max_json(Args#mrargs.start_key)},
         {end_key, maybe_replace_max_json(Args#mrargs.end_key)},
         {direction, Args#mrargs.direction},
@@ -398,6 +399,11 @@ apply_opts([{update, false} | Rest], Args) ->
         update = false
     },
     apply_opts(Rest, NewArgs);
+apply_opts([{partition, <<>>} | Rest], Args) ->
+    apply_opts(Rest, Args);
+apply_opts([{partition, Partition} | Rest], Args) when is_binary(Partition) ->
+    NewArgs = couch_mrview_util:set_extra(Args, partition, Partition),
+    apply_opts(Rest, NewArgs);
 apply_opts([{_, _} | Rest], Args) ->
     % Ignore unknown options
     apply_opts(Rest, Args).

src/mango/src/mango_error.erl

@@ -104,6 +104,13 @@ info(mango_idx, {invalid_index_type, BadType}) ->
         <<"invalid_index">>,
         fmt("Invalid type for index: ~s", [BadType])
     };
+info(mango_idx, {partitoned_option_mismatch, BadDDoc}) ->
+    {
+        400,
+        <<"invalid_partitioned_option">>,
+        fmt("Requested partitioned option does not match existing value on"
+            " design document ~s", [BadDDoc])
+    };
 info(mango_idx, invalid_query_ddoc_language) ->
     {
         400,

src/mango/src/mango_httpd.erl

@@ -170,7 +170,8 @@ handle_index_req(#httpd{path_parts=[_, _, _DDocId0, _Type, _Name]}=Req, _Db) ->
 
 handle_explain_req(#httpd{method='POST'}=Req, Db) ->
     chttpd:validate_ctype(Req, "application/json"),
-    {ok, Opts0} = mango_opts:validate_find(chttpd:json_body_obj(Req)),
+    Body = maybe_set_partition(Req),
+    {ok, Opts0} = mango_opts:validate_find(Body),
     {value, {selector, Sel}, Opts} = lists:keytake(selector, 1, Opts0),
     Resp = mango_crud:explain(Db, Sel, Opts),
     chttpd:send_json(Req, Resp);
@@ -181,7 +182,8 @@ handle_explain_req(Req, _Db) ->
 
 handle_find_req(#httpd{method='POST'}=Req, Db) ->
     chttpd:validate_ctype(Req, "application/json"),
-    {ok, Opts0} = mango_opts:validate_find(chttpd:json_body_obj(Req)),
+    Body = maybe_set_partition(Req),
+    {ok, Opts0} = mango_opts:validate_find(Body),
     {value, {selector, Sel}, Opts} = lists:keytake(selector, 1, Opts0),
     {ok, Resp0} = start_find_resp(Req),
     {ok, AccOut} = run_find(Resp0, Db, Sel, Opts),
@@ -224,6 +226,23 @@ get_idx_del_opts(Req) ->
     end.
 
 
+maybe_set_partition(Req) ->
+    {Props} = chttpd:json_body_obj(Req),
+    case chttpd:qs_value(Req, "partition", undefined) of
+        undefined ->
+            {Props};
+        Partition ->
+            case couch_util:get_value(<<"partition">>, Props) of
+                undefined ->
+                    {[{<<"partition">>, ?l2b(Partition)} | Props]};
+                Partition ->
+                    {Props};
+                OtherPartition ->
+                    ?MANGO_ERROR({bad_partition, OtherPartition})
+            end
+    end.
+
+
 convert_to_design_id(DDocId) ->
     case DDocId of
         <<"_design/", _/binary>> -> DDocId;

src/mango/src/mango_idx.erl

@@ -33,6 +33,7 @@
     name/1,
     type/1,
     def/1,
+    partitioned/1,
     opts/1,
     columns/1,
     is_usable/3,
@@ -59,8 +60,10 @@ list(Db) ->
 
 get_usable_indexes(Db, Selector, Opts) ->
     ExistingIndexes = mango_idx:list(Db),
-
-    GlobalIndexes = mango_cursor:remove_indexes_with_partial_filter_selector(ExistingIndexes),
+    MatchingPartitionIndexes = filter_partition_indexes(ExistingIndexes, Opts),
+    GlobalIndexes = mango_cursor:remove_indexes_with_partial_filter_selector(
+            MatchingPartitionIndexes
+        ),
     UserSpecifiedIndex = mango_cursor:maybe_filter_indexes_by_ddoc(ExistingIndexes, Opts),
     UsableIndexes0 = lists:usort(GlobalIndexes ++ UserSpecifiedIndex),
 
@@ -110,6 +113,7 @@ new(Db, Opts) ->
         name = IdxName,
         type = Type,
         def = Def,
+        partitioned = get_idx_partitioned(Opts),
         opts = filter_opts(Opts)
     }}.
 
@@ -121,10 +125,11 @@ validate_new(Idx, Db) ->
 
 add(DDoc, Idx) ->
     Mod = idx_mod(Idx),
-    {ok, NewDDoc} = Mod:add(DDoc, Idx),
+    {ok, NewDDoc1} = Mod:add(DDoc, Idx),
+    NewDDoc2 = set_ddoc_partitioned(NewDDoc1, Idx),
     % Round trip through JSON for normalization
-    Body = ?JSON_DECODE(?JSON_ENCODE(NewDDoc#doc.body)),
-    {ok, NewDDoc#doc{body = Body}}.
+    Body = ?JSON_DECODE(?JSON_ENCODE(NewDDoc2#doc.body)),
+    {ok, NewDDoc2#doc{body = Body}}.
 
 
 remove(DDoc, Idx) ->
@@ -176,7 +181,8 @@ from_ddoc(Db, {Props}) ->
     lists:map(fun(Idx) ->
         Idx#idx{
             dbname = DbName,
-            ddoc = DDoc
+            ddoc = DDoc,
+            partitioned = set_idx_partitioned(Db, Props)
         }
     end, Idxs).
 
@@ -213,6 +219,10 @@ def(#idx{def=Def}) ->
     Def.
 
 
+partitioned(#idx{partitioned=Partitioned}) ->
+    Partitioned.
+
+
 opts(#idx{opts=Opts}) ->
     Opts.
 
@@ -329,6 +339,87 @@ gen_name(Idx, Opts0) ->
     mango_util:enc_hex(Sha).
 
 
+get_idx_partitioned(Opts) ->
+    case proplists:get_value(partitioned, Opts) of
+        B when is_boolean(B) ->
+            B;
+        default ->
+            undefined
+    end.
+
+
+set_ddoc_partitioned(DDoc, Idx) ->
+    % We have to verify that the new index being added
+    % to this design document either matches the current
+    % ddoc's design options *or* this is a new design doc
+    #doc{
+        id = DDocId,
+        revs = Revs,
+        body = {BodyProps}
+    } = DDoc,
+    OldDOpts = couch_util:get_value(<<"options">>, BodyProps),
+    OldOpt = case OldDOpts of
+        {OldDOptProps} when is_list(OldDOptProps) ->
+            couch_util:get_value(<<"partitioned">>, OldDOptProps);
+        _ ->
+            undefined
+    end,
+    % If new matches old we're done
+    if Idx#idx.partitioned == OldOpt -> DDoc; true ->
+        % If we're creating a ddoc then we can set the options
+        case Revs == {0, []} of
+            true when Idx#idx.partitioned /= undefined ->
+                set_ddoc_partitioned_option(DDoc, Idx#idx.partitioned);
+            true when Idx#idx.partitioned == undefined ->
+                DDoc;
+            false ->
+                ?MANGO_ERROR({partitioned_option_mismatch, DDocId})
+        end
+    end.
+
+
+set_ddoc_partitioned_option(DDoc, Partitioned) ->
+    #doc{
+        body = {BodyProps}
+    } = DDoc,
+    NewProps = case couch_util:get_value(<<"options">>, BodyProps) of
+        {Existing} when is_list(Existing) ->
+            Opt = {<<"partitioned">>, Partitioned},
+            New = lists:keystore(<<"partitioned">>, 1, Existing, Opt),
+            lists:keystore(<<"options">>, 1, BodyProps, {<<"options">>, New});
+        undefined ->
+            New = {<<"options">>, {[{<<"partitioned">>, Partitioned}]}},
+            lists:keystore(<<"options">>, 1, BodyProps, New)
+    end,
+    DDoc#doc{body = {NewProps}}.
+
+
+set_idx_partitioned(Db, DDocProps) ->
+    Default = fabric_util:is_partitioned(Db),
+    case couch_util:get_value(<<"options">>, DDocProps) of
+        {DesignOpts} ->
+            case couch_util:get_value(<<"partitioned">>, DesignOpts) of
+                P when is_boolean(P) ->
+                    P;
+                undefined ->
+                    Default
+            end;
+        undefined ->
+            Default
+    end.
+
+
+filter_partition_indexes(Indexes, Opts) ->
+    PFilt = case couch_util:get_value(partition, Opts) of
+        <<>> ->
+            fun(#idx{partitioned = P}) -> not P end;
+        Partition when is_binary(Partition) ->
+            fun(#idx{partitioned = P}) -> P end
+    end,
+    Filt = fun(Idx) -> type(Idx) == <<"special">> orelse PFilt(Idx) end,
+    lists:filter(Filt, Indexes).
+
+
 filter_opts([]) ->
     [];
 filter_opts([{user_ctx, _} | Rest]) ->
@@ -341,6 +432,8 @@ filter_opts([{type, _} | Rest]) ->
     filter_opts(Rest);
 filter_opts([{w, _} | Rest]) ->
     filter_opts(Rest);
+filter_opts([{partitioned, _} | Rest]) ->
+    filter_opts(Rest);
 filter_opts([Opt | Rest]) ->
     [Opt | filter_opts(Rest)].
 

src/mango/src/mango_idx.hrl

@@ -16,5 +16,6 @@
     name,
     type,
     def,
+    partitioned,
     opts
 }).

src/mango/src/mango_opts.erl

@@ -34,6 +34,7 @@
     validate_sort/1,
     validate_fields/1,
     validate_bulk_delete/1,
+    validate_partitioned/1,
 
     default_limit/0
 ]).
@@ -70,6 +71,12 @@ validate_idx_create({Props}) ->
             {optional, true},
             {default, 2},
             {validator, fun is_pos_integer/1}
+        ]},
+        {<<"partitioned">>, [
+            {tag, partitioned},
+            {optional, true},
+            {default, default},
+            {validator, fun validate_partitioned/1}
         ]}
     ],
     validate(Props, Opts).
@@ -117,6 +124,12 @@ validate_find({Props}) ->
             {default, []},
             {validator, fun validate_fields/1}
         ]},
+        {<<"partition">>, [
+            {tag, partition},
+            {optional, true},
+            {default, <<>>},
+            {validator, fun validate_partition/1}
+        ]},
         {<<"r">>, [
             {tag, r},
             {optional, true},
@@ -296,6 +309,23 @@ validate_fields(Value) ->
     mango_fields:new(Value).
 
 
+validate_partitioned(true) ->
+    {ok, true};
+validate_partitioned(false) ->
+    {ok, false};
+validate_partitioned(default) ->
+    {ok, default};
+validate_partitioned(Else) ->
+    ?MANGO_ERROR({invalid_partitioned_value, Else}).
+
+
+validate_partition(<<>>) ->
+    {ok, <<>>};
+validate_partition(Partition) ->
+    couch_partition:validate_partition(Partition),
+    {ok, Partition}.
+
+
 validate_opts([], Props, Acc) ->
     {Props, lists:reverse(Acc)};
 validate_opts([{Name, Desc} | Rest], Props, Acc) ->