Implement _all_docs and _find support

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

13 files changed, 175 insertions, 36 deletions

diff --git a/src/chttpd/src/chttpd_db.erl b/src/chttpd/src/chttpd_db.erl
index bd64ad04f..a68ea8cde 100644
--- a/src/chttpd/src/chttpd_db.erl
+++ b/src/chttpd/src/chttpd_db.erl
@@ -715,15 +715,18 @@ multi_all_docs_view(Req, Db, OP, Queries) ->
     {ok, Resp1} = chttpd:send_delayed_chunk(VAcc2#vacc.resp, "\r\n]}"),
     chttpd:end_delayed_json_response(Resp1).
 
-all_docs_view(Req, Db, Keys, OP) ->
+all_docs_view(#httpd{path_parts=[DbName | _]}=Req, Db, Keys, OP) ->
+    Partitioned = mem3:is_partitioned(DbName),
     Args0 = couch_mrview_http:parse_params(Req, Keys),
     Args1 = Args0#mrargs{view_type=map},
-    Args2 = couch_mrview_util:validate_args(Args1),
-    Args3 = set_namespace(OP, Args2),
+    Args2 = couch_mrview_util:set_extra(Args1, style, all_docs),
+    Args3 = couch_mrview_util:set_extra(Args2, partitioned, Partitioned),
+    Args4 = couch_mrview_util:validate_args(Args3),
+    Args5 = set_namespace(OP, Args4),
     Options = [{user_ctx, Req#httpd.user_ctx}],
     Max = chttpd:chunked_response_buffer_size(),
     VAcc = #vacc{db=Db, req=Req, threshold=Max},
-    {ok, Resp} = fabric:all_docs(Db, Options, fun couch_mrview_http:view_cb/2, VAcc, Args3),
+    {ok, Resp} = fabric:all_docs(Db, Options, fun couch_mrview_http:view_cb/2, VAcc, Args5),
     {ok, Resp#vacc.resp}.
 
 db_doc_req(#httpd{method='DELETE'}=Req, Db, DocId) ->
diff --git a/src/couch_mrview/src/couch_mrview.erl b/src/couch_mrview/src/couch_mrview.erl
index 09945f555..f5963e70f 100644
--- a/src/couch_mrview/src/couch_mrview.erl
+++ b/src/couch_mrview/src/couch_mrview.erl
@@ -228,12 +228,13 @@ query_all_docs(Db, Args0, Callback, Acc) ->
         couch_index_util:hexsig(couch_hash:md5_hash(term_to_binary(Info)))
     end),
     Args1 = Args0#mrargs{view_type=map},
-    Args2 = couch_mrview_util:validate_args(Args1),
-    {ok, Acc1} = case Args2#mrargs.preflight_fun of
+    Args2 = couch_mrview_util:set_extra(Args1, style, all_docs),
+    Args3 = couch_mrview_util:validate_args(Args2),
+    {ok, Acc1} = case Args3#mrargs.preflight_fun of
         PFFun when is_function(PFFun, 2) -> PFFun(Sig, Acc);
         _ -> {ok, Acc}
     end,
-    all_docs_fold(Db, Args2, Callback, Acc1).
+    all_docs_fold(Db, Args3, Callback, Acc1).
 
 
 query_view(Db, DDoc, VName) ->
diff --git a/src/couch_mrview/src/couch_mrview_updater.erl b/src/couch_mrview/src/couch_mrview_updater.erl
index d02bf9707..7870d72c9 100644
--- a/src/couch_mrview/src/couch_mrview_updater.erl
+++ b/src/couch_mrview/src/couch_mrview_updater.erl
@@ -319,7 +319,9 @@ write_kvs(State, UpdateSeq, ViewKVs, DocIdKeys, Seqs, Log0) ->
         design_opts=DesignOpts
     } = State,
 
-    Partitioned = couch_util:get_value(<<"partitioned">>, DesignOpts, false),
+    DbPartitioned = mem3:is_partitioned(State#mrst.db_name),
+    Partitioned = couch_util:get_value(<<"partitioned">>, DesignOpts, DbPartitioned),
+
     Revs = dict:from_list(dict:fetch_keys(Log0)),
 
     Log = dict:fold(fun({Id, _Rev}, DIKeys, Acc) ->
diff --git a/src/couch_mrview/src/couch_mrview_util.erl b/src/couch_mrview/src/couch_mrview_util.erl
index 733b0c08a..fcdadda5f 100644
--- a/src/couch_mrview/src/couch_mrview_util.erl
+++ b/src/couch_mrview/src/couch_mrview_util.erl
@@ -41,8 +41,11 @@
 -define(GET_VIEW_RETRY_COUNT, 1).
 -define(GET_VIEW_RETRY_DELAY, 50).
 -define(LOWEST_KEY, null).
--define(HIGHEST_KEY, {[{<<239, 191, 176>>, null}]}). % is {"\ufff0": null}
-
+-define(HIGHEST_KEY, {<<255, 255, 255, 255>>}).
+-define(PARTITION_START(P), <<P/binary, $:>>).
+-define(PARTITION_END(P), <<P/binary, $;>>).
+-define(LOWEST(A, B), (if A < B -> A; true -> B end)).
+-define(HIGHEST(A, B), (if A > B -> A; true -> B end)).
 
 -include_lib("couch/include/couch_db.hrl").
 -include_lib("couch_mrview/include/couch_mrview.hrl").
@@ -618,20 +621,38 @@ validate_args(Args) ->
         _ -> mrverror(<<"Invalid value for `sorted`.">>)
     end,
 
-    case {get_extra(Args, partitioned, false), get_extra(Args, partition)} of
-        {true, undefined} ->
+    Style = get_extra(Args, style, normal),
+    Partitioned = get_extra(Args, partitioned, false),
+    Partition = get_extra(Args, partition),
+
+    case {Style, Partitioned, Partition} of
+        {all_docs, true, _} ->
+            ok; % _all_docs can be called with or without partition parameter.
+        {all_docs, false, undefined} ->
+            ok;
+        {all_docs, false, _Partition} ->
+            mrverror(<<"`partition` parameter is not supported in this db.">>);
+        {normal, true, undefined} ->
             mrverror(<<"`partition` parameter is mandatory for queries to this view.">>);
-        {true, Partition} ->
+        {normal, true, Partition} ->
             couch_doc:validate_docid(Partition);
-        {false, undefined} ->
+        {normal, false, undefined} ->
             ok;
-        {false, _Partition} ->
+        {normal, false, _Partition} ->
             mrverror(<<"`partition` parameter is not supported in this view.">>)
     end,
 
-    Args1 = case get_extra(Args, partitioned, false) of
-        true  -> apply_partition(Args);
-        false -> Args
+    Args1 = case {Style, Partitioned, Partition} of
+        {all_docs, true, undefined} ->
+            Args;
+        {all_docs, true, Partition} ->
+            apply_partition(Args, all_docs);
+        {all_docs, false, _} ->
+            Args;
+        {normal, true, _} ->
+            apply_partition(Args, normal);
+        {normal, false, _} ->
+            Args
     end,
 
     Args1#mrargs{
@@ -652,9 +673,12 @@ determine_group_level(#mrargs{group=true, group_level=undefined}) ->
 determine_group_level(#mrargs{group_level=GroupLevel}) ->
     GroupLevel.
 
-apply_partition(#mrargs{} = Args0) ->
+apply_partition(#mrargs{} = Args0, Style) ->
     Partition = get_extra(Args0, partition),
-    apply_partition(Partition, Args0).
+    case Style of
+        normal   -> apply_partition(Partition, Args0);
+        all_docs -> apply_all_docs_partition(Partition, Args0)
+    end;
 
 apply_partition(_Partition, #mrargs{keys=[{p, _, _} | _]} = Args) ->
     Args; % already applied
@@ -685,6 +709,33 @@ apply_partition(Partition, Args) ->
         end_key = {p, Partition, EK0}
     }.
 
+%% all_docs is special as it's not really a view and is already
+%% effectively partitioned as the partition is a prefix of all keys.
+apply_all_docs_partition(Partition, #mrargs{direction=fwd, start_key=undefined, end_key=undefined} = Args) ->
+    Args#mrargs{start_key = ?PARTITION_START(Partition), end_key = ?PARTITION_END(Partition)};
+
+apply_all_docs_partition(Partition, #mrargs{direction=rev, start_key=undefined, end_key=undefined} = Args) ->
+    Args#mrargs{start_key = ?PARTITION_END(Partition), end_key = ?PARTITION_START(Partition)};
+
+apply_all_docs_partition(Partition, #mrargs{direction=fwd, start_key=SK, end_key=undefined} = Args) ->
+    Args#mrargs{start_key = ?HIGHEST(?PARTITION_START(Partition), SK), end_key = ?PARTITION_END(Partition)};
+
+apply_all_docs_partition(Partition, #mrargs{direction=rev, start_key=SK, end_key=undefined} = Args) ->
+    Args#mrargs{start_key = ?LOWEST(?PARTITION_END(Partition), SK), end_key = ?PARTITION_START(Partition)};
+
+apply_all_docs_partition(Partition, #mrargs{direction=fwd, start_key=undefined, end_key=EK} = Args) ->
+    Args#mrargs{start_key = ?PARTITION_START(Partition), end_key = ?LOWEST(?PARTITION_END(Partition), EK)};
+
+apply_all_docs_partition(Partition, #mrargs{direction=rev, start_key=undefined, end_key=EK} = Args) ->
+    Args#mrargs{start_key = ?PARTITION_END(Partition), end_key = ?HIGHEST(?PARTITION_START(Partition), EK)};
+
+apply_all_docs_partition(Partition, #mrargs{direction=fwd, start_key=SK, end_key=EK} = Args) ->
+    Args#mrargs{start_key = ?HIGHEST(?PARTITION_START(Partition), SK), end_key = ?LOWEST(?PARTITION_END(Partition), EK)};
+
+apply_all_docs_partition(Partition, #mrargs{direction=rev, start_key=SK, end_key=EK} = Args) ->
+    Args#mrargs{start_key = ?LOWEST(?PARTITION_END(Partition), SK), end_key = ?HIGHEST(?PARTITION_START(Partition), EK)}.
+
+
 check_range(#mrargs{start_key=undefined}, _Cmp) ->
     ok;
 check_range(#mrargs{end_key=undefined}, _Cmp) ->
diff --git a/src/fabric/src/fabric_view.erl b/src/fabric/src/fabric_view.erl
index 844b44dfd..c0e29740d 100644
--- a/src/fabric/src/fabric_view.erl
+++ b/src/fabric/src/fabric_view.erl
@@ -122,7 +122,7 @@ maybe_send_row(State) ->
         user_acc = AccIn,
         query_args = QueryArgs
     } = State,
-    Partitioned = couch_mrview_util:get_extra(QueryArgs, partitioned, false),
+    Partitioned = couch_mrview_util:get_extra(QueryArgs, partitioned),
     case fabric_dict:any(0, Counters) of
     true ->
         {ok, State};
diff --git a/src/fabric/src/fabric_view_all_docs.erl b/src/fabric/src/fabric_view_all_docs.erl
index ac16dac52..d515ab830 100644
--- a/src/fabric/src/fabric_view_all_docs.erl
+++ b/src/fabric/src/fabric_view_all_docs.erl
@@ -118,7 +118,7 @@ go(DbName, _Options, Workers, QueryArgs, Callback, Acc0) ->
     #mrargs{limit = Limit, skip = Skip, update_seq = UpdateSeq} = QueryArgs,
     State = #collector{
         db_name = DbName,
-        query_args = QueryArgs,
+        query_args = couch_mrview_util:set_extra(QueryArgs, style, all_docs),
         callback = Callback,
         counters = fabric_dict:init(Workers, 0),
         skip = Skip,
diff --git a/src/mango/src/mango_cursor_view.erl b/src/mango/src/mango_cursor_view.erl
index 174381e4a..1ac10b348 100644
--- a/src/mango/src/mango_cursor_view.erl
+++ b/src/mango/src/mango_cursor_view.erl
@@ -32,6 +32,7 @@
 -include_lib("fabric/include/fabric.hrl").
 
 -include("mango_cursor.hrl").
+-include("mango_idx.hrl").
 -include("mango_idx_view.hrl").
 
 -define(HEARTBEAT_INTERVAL_IN_USEC, 4000000).
@@ -76,7 +77,8 @@ explain(Cursor) ->
         {direction, Args#mrargs.direction},
         {stable, Args#mrargs.stable},
         {update, Args#mrargs.update},
-        {conflicts, Args#mrargs.conflicts}
+        {conflicts, Args#mrargs.conflicts},
+        {partition, couch_mrview_util:get_extra(Args, partition, null)}
     ]}}].
 
 
@@ -98,15 +100,29 @@ maybe_replace_max_json([H | T] = EndKey) when is_list(EndKey) ->
 maybe_replace_max_json(EndKey) ->
     EndKey.
 
-base_args(#cursor{index = Idx, selector = Selector} = Cursor) ->
-    #mrargs{
+base_args(#cursor{index = Idx, opts = Opts, selector = Selector} = Cursor) ->
+    Args1 = #mrargs{
         view_type = map,
         reduce = false,
         start_key = mango_idx:start_key(Idx, Cursor#cursor.ranges),
         end_key = mango_idx:end_key(Idx, Cursor#cursor.ranges),
         include_docs = true,
         extra = [{callback, {?MODULE, view_cb}}, {selector, Selector}]
-    }.
+    },
+    Partitioned = couch_util:get_value(partitioned, Idx#idx.design_opts),
+    Args2 = couch_mrview_util:set_extra(Args1, partitioned, Partitioned),
+    Args3 = case couch_util:get_value(partition, Opts) of
+        <<>> ->
+            Args2;
+        Partition ->
+            couch_mrview_util:set_extra(Args2, partition, Partition)
+    end,
+    add_style(Idx, Args3).
+
+add_style(#idx{def = all_docs}, Args) ->
+    couch_mrview_util:set_extra(Args, style, all_docs);
+add_style(_, Args) ->
+    Args.
 
 
 execute(#cursor{db = Db, index = Idx, execution_stats = Stats} = Cursor0, UserFun, UserAcc) ->
diff --git a/src/mango/src/mango_httpd.erl b/src/mango/src/mango_httpd.erl
index 2e8777135..cac8e3d04 100644
--- a/src/mango/src/mango_httpd.erl
+++ b/src/mango/src/mango_httpd.erl
@@ -81,9 +81,9 @@ handle_index_req(#httpd{method='GET', path_parts=[_, _]}=Req, Db) ->
     JsonIdxs = lists:sublist(JsonIdxs0, Skip+1, Limit),
 	chttpd:send_json(Req, {[{total_rows, TotalRows}, {indexes, JsonIdxs}]});
 
-handle_index_req(#httpd{method='POST', path_parts=[_, _]}=Req, Db) ->
+handle_index_req(#httpd{method='POST', path_parts=[DbName, _]}=Req, Db) ->
     chttpd:validate_ctype(Req, "application/json"),
-    {ok, Opts} = mango_opts:validate_idx_create(chttpd:json_body_obj(Req)),
+    {ok, Opts} = mango_opts:validate_idx_create(DbName, chttpd:json_body_obj(Req)),
     {ok, Idx0} = mango_idx:new(Db, Opts),
     {ok, Idx} = mango_idx:validate_new(Idx0, Db),
     DbOpts = [{user_ctx, Req#httpd.user_ctx}, deleted, ejson_body],
diff --git a/src/mango/src/mango_idx.erl b/src/mango/src/mango_idx.erl
index 8af92b946..e051218f4 100644
--- a/src/mango/src/mango_idx.erl
+++ b/src/mango/src/mango_idx.erl
@@ -58,8 +58,9 @@ list(Db) ->
 
 
 get_usable_indexes(Db, Selector, Opts) ->
-    ExistingIndexes = mango_idx:list(Db),
-
+    PQ = is_partitioned_query(Opts),
+    ExistingIndexes = filter_indexes_by_partitioned(
+        mango_idx:list(Db), PQ),
     GlobalIndexes = mango_cursor:remove_indexes_with_partial_filter_selector(ExistingIndexes),
     UserSpecifiedIndex = mango_cursor:maybe_filter_indexes_by_ddoc(ExistingIndexes, Opts),
     UsableIndexes0 = lists:usort(GlobalIndexes ++ UserSpecifiedIndex),
@@ -68,13 +69,36 @@ get_usable_indexes(Db, Selector, Opts) ->
     UsableFilter = fun(I) -> is_usable(I, Selector, SortFields) end,
 
     case lists:filter(UsableFilter, UsableIndexes0) of
-        [] -> 
+        [] ->
             ?MANGO_ERROR({no_usable_index, missing_sort_index});
-        UsableIndexes -> 
+        UsableIndexes ->
             UsableIndexes
     end.
 
 
+filter_indexes_by_partitioned(Indexes, PQ) ->
+    filter_indexes_by_partitioned(Indexes, PQ, []).
+
+
+filter_indexes_by_partitioned([], _PQ, Acc) ->
+    lists:reverse(Acc);
+filter_indexes_by_partitioned([Idx | Rest], PQ, Acc) ->
+    {partitioned, PI} = lists:keyfind(partitioned, 1, Idx#idx.design_opts),
+    case {Idx#idx.def, PI, PQ} of
+        {all_docs, _, _} ->
+            % all_docs works both ways.
+            filter_indexes_by_partitioned(Rest, PQ, [Idx | Acc]);
+        {_, Same, Same} ->
+            filter_indexes_by_partitioned(Rest, PQ, [Idx | Acc]);
+        {_, _, _} ->
+             filter_indexes_by_partitioned(Rest, PQ, Acc)
+    end.
+
+
+is_partitioned_query(Opts) ->
+    lists:keyfind(partition, 1, Opts) /= {partition, <<>>}.
+
+
 recover(Db) ->
     {ok, DDocs0} = mango_util:open_ddocs(Db),
     Pred = fun({Props}) ->
@@ -101,6 +125,7 @@ get_sort_fields(Opts) ->
 
 new(Db, Opts) ->
     Def = get_idx_def(Opts),
+    DesignOpts = get_idx_design_opts(Db, Opts),
     Type = get_idx_type(Opts),
     IdxName = get_idx_name(Def, Opts),
     DDoc = get_idx_ddoc(Def, Opts),
@@ -110,6 +135,7 @@ new(Db, Opts) ->
         name = IdxName,
         type = Type,
         def = Def,
+        design_opts = DesignOpts,
         opts = filter_opts(Opts)
     }}.
 
@@ -182,11 +208,13 @@ from_ddoc(Db, {Props}) ->
 
 
 special(Db) ->
+    Partitioned = mem3:is_partitioned(Db),
     AllDocs = #idx{
         dbname = db_to_name(Db),
         name = <<"_all_docs">>,
         type = <<"special">>,
         def = all_docs,
+        design_opts = [{partitioned, Partitioned}],
         opts = []
     },
     % Add one for _update_seq
@@ -285,6 +313,12 @@ get_idx_def(Opts) ->
     end.
 
 
+get_idx_design_opts(Db, Opts) ->
+    DbPartitioned = mem3:is_partitioned(couch_db:name(Db)),
+    Partitioned = proplists:get_value(partitioned, Opts, DbPartitioned),
+    [{partitioned, Partitioned}].
+
+
 get_idx_type(Opts) ->
     case proplists:get_value(type, Opts) of
         <<"json">> -> <<"json">>;
@@ -341,6 +375,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)].
 
diff --git a/src/mango/src/mango_idx.hrl b/src/mango/src/mango_idx.hrl
index 712031b75..179911f46 100644
--- a/src/mango/src/mango_idx.hrl
+++ b/src/mango/src/mango_idx.hrl
@@ -16,5 +16,6 @@
     name,
     type,
     def,
+    design_opts,
     opts
 }).
diff --git a/src/mango/src/mango_idx_view.erl b/src/mango/src/mango_idx_view.erl
index 2d784b638..ef1ca5912 100644
--- a/src/mango/src/mango_idx_view.erl
+++ b/src/mango/src/mango_idx_view.erl
@@ -54,7 +54,8 @@ add(#doc{body={Props0}}=DDoc, Idx) ->
     NewView = make_view(Idx),
     Views2 = lists:keystore(element(1, NewView), 1, Views1, NewView),
     Props1 = lists:keystore(<<"views">>, 1, Props0, {<<"views">>, {Views2}}),
-    {ok, DDoc#doc{body={Props1}}}.
+    Props2 = lists:keystore(<<"options">>, 1, Props1, {<<"options">>, {Idx#idx.design_opts}}),
+    {ok, DDoc#doc{body={Props2}}}.
 
 
 remove(#doc{body={Props0}}=DDoc, Idx) ->
@@ -78,6 +79,7 @@ remove(#doc{body={Props0}}=DDoc, Idx) ->
 
 
 from_ddoc({Props}) ->
+    DesignOpts = validate_design_opts(Props),
     case lists:keyfind(<<"views">>, 1, Props) of
         {<<"views">>, {Views}} when is_list(Views) ->
             lists:flatmap(fun({Name, {VProps}}) ->
@@ -89,6 +91,7 @@ from_ddoc({Props}) ->
                         type = <<"json">>,
                         name = Name,
                         def = Def,
+                        design_opts = DesignOpts,
                         opts = Opts
                         },
                         [I]
@@ -104,6 +107,7 @@ to_json(Idx) ->
         {ddoc, Idx#idx.ddoc},
         {name, Idx#idx.name},
         {type, Idx#idx.type},
+        {design_opts, {Idx#idx.design_opts}},
         {def, {def_to_json(Idx#idx.def)}}
     ]}.
 
@@ -248,6 +252,16 @@ validate_ddoc(VProps) ->
     end.
 
 
+validate_design_opts(Props) ->
+    case lists:keyfind(<<"options">>, 1, Props) of
+        {<<"options">>, {[{<<"partitioned">>, P}]}}
+          when is_boolean(P) ->
+            [{partitioned, P}];
+        _ ->
+            []
+    end.
+
+
 % This function returns a list of indexes that
 % can be used to restrict this query. This works by
 % searching the selector looking for field names that
diff --git a/src/mango/src/mango_opts.erl b/src/mango/src/mango_opts.erl
index 7bae9c90d..d846287c3 100644
--- a/src/mango/src/mango_opts.erl
+++ b/src/mango/src/mango_opts.erl
@@ -13,7 +13,7 @@
 -module(mango_opts).
 
 -export([
-    validate_idx_create/1,
+    validate_idx_create/2,
     validate_find/1
 ]).
 
@@ -42,11 +42,17 @@
 -include("mango.hrl").
 
 
-validate_idx_create({Props}) ->
+validate_idx_create(DbName, {Props}) ->
     Opts = [
         {<<"index">>, [
             {tag, def}
         ]},
+        {<<"partitioned">>, [
+            {tag, partitioned},
+            {optional, true},
+            {default, mem3:is_partitioned(DbName)},
+            {validator, fun is_boolean/1}
+        ]},
         {<<"type">>, [
             {tag, type},
             {optional, true},
@@ -81,6 +87,12 @@ validate_find({Props}) ->
             {tag, selector},
             {validator, fun validate_selector/1}
         ]},
+        {<<"partition">>, [
+            {tag, partition},
+            {optional, true},
+            {default, <<>>},
+            {validator, fun is_string/1}
+        ]},
         {<<"use_index">>, [
             {tag, use_index},
             {optional, true},
diff --git a/src/mem3/src/mem3.erl b/src/mem3/src/mem3.erl
index 0ad916403..d6d14dde3 100644
--- a/src/mem3/src/mem3.erl
+++ b/src/mem3/src/mem3.erl
@@ -372,7 +372,10 @@ is_partitioned(DbName0) when is_binary(DbName0) ->
             false
     end;
 
-is_partitioned(Shards) when is_list(Shards) ->
+is_partitioned([#shard{} | _] = Shards) ->
+    lists:all(fun is_partitioned/1, Shards);
+
+is_partitioned([#ordered_shard{} | _] = Shards) ->
     lists:all(fun is_partitioned/1, Shards);
 
 is_partitioned(#shard{opts=Opts}) ->