add local_docs to fold_doc with docids

4 files changed, 292 insertions, 23 deletions

diff --git a/src/fabric/src/fabric2_db.erl b/src/fabric/src/fabric2_db.erl
index 740f9abf6..8764d4e18 100644
--- a/src/fabric/src/fabric2_db.erl
+++ b/src/fabric/src/fabric2_db.erl
@@ -976,16 +976,6 @@ fold_docs(Db, DocIds, UserFun, UserAcc0, Options) ->
             NeedsTreeOpts = [revs_info, conflicts, deleted_conflicts],
             NeedsTree = (Options -- NeedsTreeOpts /= Options),
 
-            FetchRevs = case NeedsTree of
-                true ->
-                    fun(DocId) ->
-                        fabric2_fdb:get_all_revs_future(TxDb, DocId)
-                    end;
-                false ->
-                    fun(DocId) ->
-                        fabric2_fdb:get_winning_revs_future(TxDb, DocId, 1)
-                    end
-            end,
             InitAcc = #{
                 revs_q => queue:new(),
                 revs_count => 0,
@@ -1001,7 +991,7 @@ fold_docs(Db, DocIds, UserFun, UserAcc0, Options) ->
                     revs_q := RevsQ,
                     revs_count := RevsCount
                 } = Acc,
-                Future = FetchRevs(DocId),
+                Future = fold_docs_get_revs(TxDb, DocId, NeedsTree),
                 NewAcc = Acc#{
                     revs_q := queue:in({DocId, Future}, RevsQ),
                     revs_count := RevsCount + 1
@@ -1262,6 +1252,47 @@ drain_all_deleted_info_futures(FutureQ, UserFun, Acc) ->
     end.
 
 
+fold_docs_get_revs(Db, <<?LOCAL_DOC_PREFIX, _/binary>> = DocId, _) ->
+    fabric2_fdb:get_local_doc_rev_future(Db, DocId);
+
+fold_docs_get_revs(Db, DocId, true) ->
+    fabric2_fdb:get_all_revs_future(Db, DocId);
+
+fold_docs_get_revs(Db, DocId, false) ->
+    fabric2_fdb:get_winning_revs_future(Db, DocId, 1).
+
+
+fold_docs_get_revs_wait(_Db, <<?LOCAL_DOC_PREFIX, _/binary>>, RevsFuture) ->
+    Rev = fabric2_fdb:get_local_doc_rev_wait(RevsFuture),
+    [Rev];
+
+fold_docs_get_revs_wait(Db, _DocId, RevsFuture) ->
+    fabric2_fdb:get_revs_wait(Db, RevsFuture).
+
+
+fold_docs_get_doc_body_future(Db, <<?LOCAL_DOC_PREFIX, _/binary>> = DocId,
+        [Rev]) ->
+    fabric2_fdb:get_local_doc_body_future(Db, DocId, Rev);
+
+fold_docs_get_doc_body_future(Db, DocId, Revs) ->
+    Winner = get_rev_winner(Revs),
+    fabric2_fdb:get_doc_body_future(Db, DocId, Winner).
+
+
+fold_docs_get_doc_body_wait(Db, <<?LOCAL_DOC_PREFIX, _/binary>> = DocId, [Rev],
+        _DocOpts, BodyFuture) ->
+    case fabric2_fdb:get_local_doc_body_wait(Db, DocId, Rev, BodyFuture) of
+        {not_found, missing} -> {not_found, missing};
+        Doc -> {ok, Doc}
+    end;
+
+fold_docs_get_doc_body_wait(Db, DocId, Revs, DocOpts, BodyFuture) ->
+    RevInfo = get_rev_winner(Revs),
+    Base = fabric2_fdb:get_doc_body_wait(Db, DocId, RevInfo,
+        BodyFuture),
+    apply_open_doc_opts(Base, Revs, DocOpts).
+
+
 drain_fold_docs_revs_futures(_TxDb, #{revs_count := C} = Acc) when C < 100 ->
     Acc;
 drain_fold_docs_revs_futures(TxDb, Acc) ->
@@ -1284,13 +1315,12 @@ drain_one_fold_docs_revs_future(TxDb, Acc) ->
     } = Acc,
     {{value, {DocId, RevsFuture}}, RestRevsQ} = queue:out(RevsQ),
 
-    Revs = fabric2_fdb:get_revs_wait(TxDb, RevsFuture),
+    Revs = fold_docs_get_revs_wait(TxDb, DocId, RevsFuture),
     DocFuture = case Revs of
         [] ->
             {DocId, [], not_found};
         [_ | _] ->
-            Winner = get_rev_winner(Revs),
-            BodyFuture = fabric2_fdb:get_doc_body_future(TxDb, DocId, Winner),
+            BodyFuture = fold_docs_get_doc_body_future(TxDb, DocId, Revs),
             {DocId, Revs, BodyFuture}
     end,
     NewAcc = Acc#{
@@ -1328,10 +1358,7 @@ drain_one_fold_docs_body_future(TxDb, Acc) ->
         not_found ->
             {not_found, missing};
         _ ->
-            RevInfo = get_rev_winner(Revs),
-            Base = fabric2_fdb:get_doc_body_wait(TxDb, DocId, RevInfo,
-                BodyFuture),
-            apply_open_doc_opts(Base, Revs, DocOpts)
+            fold_docs_get_doc_body_wait(TxDb, DocId, Revs, DocOpts, BodyFuture)
     end,
     NewUserAcc = maybe_stop(UserFun(DocId, Doc, UserAcc)),
     Acc#{
diff --git a/src/fabric/src/fabric2_fdb.erl b/src/fabric/src/fabric2_fdb.erl
index ba57e646d..8264e8a60 100644
--- a/src/fabric/src/fabric2_fdb.erl
+++ b/src/fabric/src/fabric2_fdb.erl
@@ -51,6 +51,11 @@
     get_doc_body/3,
     get_doc_body_future/3,
     get_doc_body_wait/4,
+
+    get_local_doc_rev_future/2,
+    get_local_doc_rev_wait/1,
+    get_local_doc_body_future/3,
+    get_local_doc_body_wait/4,
     get_local_doc/2,
     get_local_doc_rev/3,
 
@@ -712,22 +717,45 @@ get_doc_body_wait(#{} = Db0, DocId, RevInfo, Future) ->
     fdb_to_doc(Db, DocId, RevPos, [Rev | RevPath], BodyRows).
 
 
-get_local_doc(#{} = Db0, <<?LOCAL_DOC_PREFIX, _/binary>> = DocId) ->
+get_local_doc_rev_future(Db, DocId) ->
     #{
         tx := Tx,
         db_prefix := DbPrefix
-    } = Db = ensure_current(Db0),
+    } = ensure_current(Db),
 
     Key = erlfdb_tuple:pack({?DB_LOCAL_DOCS, DocId}, DbPrefix),
-    Rev = erlfdb:wait(erlfdb:get(Tx, Key)),
+    erlfdb:get(Tx, Key).
+
+
+get_local_doc_rev_wait(Future) ->
+    erlfdb:wait(Future).
+
+
+get_local_doc_body_future(#{} = Db, DocId, Rev) ->
+    #{
+        tx := Tx,
+        db_prefix := DbPrefix
+    } = ensure_current(Db),
 
     Prefix = erlfdb_tuple:pack({?DB_LOCAL_DOC_BODIES, DocId}, DbPrefix),
-    Future = erlfdb:get_range_startswith(Tx, Prefix),
-    {_, Chunks} = lists:unzip(aegis:decrypt(Db, erlfdb:wait(Future))),
+    erlfdb:get_range_startswith(Tx, Prefix).
+
+
+get_local_doc_body_wait(#{} = Db0, DocId, Rev, Future) ->
+    Db = ensure_current(Db0),
 
+    {_, Chunks} = lists:unzip(aegis:decrypt(Db, erlfdb:wait(Future))),
     fdb_to_local_doc(Db, DocId, Rev, Chunks).
 
 
+get_local_doc(#{} = Db, <<?LOCAL_DOC_PREFIX, _/binary>> = DocId) ->
+    RevFuture = get_local_doc_rev_future(Db, DocId),
+    Rev = get_local_doc_rev_wait(RevFuture),
+
+    BodyFuture = get_local_doc_body_future(Db, DocId, Rev),
+    get_local_doc_body_wait(Db, DocId, Rev, BodyFuture).
+
+
 get_local_doc_rev(_Db0, <<?LOCAL_DOC_PREFIX, _/binary>> = DocId, Val) ->
     case Val of
         <<255, RevBin/binary>> ->
diff --git a/src/fabric/test/fabric2_db_fold_doc_docids_tests.erl b/src/fabric/test/fabric2_db_fold_doc_docids_tests.erl
new file mode 100644
index 000000000..b55da5363
--- /dev/null
+++ b/src/fabric/test/fabric2_db_fold_doc_docids_tests.erl
@@ -0,0 +1,150 @@
+% Licensed under the Apache License, Version 2.0 (the "License"); you may not
+% use this file except in compliance with the License. You may obtain a copy of
+% the License at
+%
+%   http://www.apache.org/licenses/LICENSE-2.0
+%
+% Unless required by applicable law or agreed to in writing, software
+% distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+% WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+% License for the specific language governing permissions and limitations under
+% the License.
+
+-module(fabric2_db_fold_doc_docids_tests).
+
+
+-include_lib("couch/include/couch_db.hrl").
+-include_lib("couch/include/couch_eunit.hrl").
+-include_lib("eunit/include/eunit.hrl").
+-include("fabric2_test.hrl").
+
+doc_fold_test_() ->
+    {
+        "Test document fold operations",
+        {
+            setup,
+            fun setup_all/0,
+            fun teardown_all/1,
+            {
+                foreach,
+                fun setup/0,
+                fun cleanup/1,
+                [
+                    ?TDEF_FE(fold_docs_simple),
+                    ?TDEF_FE(fold_docs_lots),
+                    ?TDEF_FE(fold_docs_local),
+                    ?TDEF_FE(fold_docs_mixed)
+]
+            }
+        }
+    }.
+
+
+setup_all() ->
+    test_util:start_couch([fabric]).
+
+
+teardown_all(Ctx) ->
+    test_util:stop_couch(Ctx).
+
+
+setup() ->
+    {ok, Db} = fabric2_db:create(?tempdb(), [{user_ctx, ?ADMIN_USER}]),
+    Db.
+
+
+cleanup(Db) ->
+    ok = fabric2_db:delete(fabric2_db:name(Db), []).
+
+
+fold_docs_simple(Db) ->
+    Docs = create_docs(Db, 10),
+    run_fold(Db, Docs).
+
+
+fold_docs_lots(Db) ->
+    Docs = create_docs(Db, 110),
+    run_fold(Db, Docs).
+
+
+fold_docs_local(Db) ->
+    Docs = create_local_docs(Db, 10),
+    run_fold(Db, Docs).
+
+
+fold_docs_mixed(Db) ->
+    Docs = create_mixed_docs(Db, 200),
+    run_fold(Db, Docs).
+
+
+run_fold(Db, Docs) ->
+    SortedIds = get_ids(Docs),
+    Ids = shuffle(SortedIds),
+    Returned = fabric2_fdb:transactional(Db, fun (TxDb) ->
+        fold_docs_return_ids(TxDb, Ids)
+    end),
+    ?assertEqual(Returned, Ids).
+
+
+fold_docs_return_ids(TxDb, Ids) ->
+    CB = fun(DocId, _Doc, Acc) ->
+        {ok, Acc ++ [DocId]}
+    end,
+    {ok, Acc} = fabric2_db:fold_docs(TxDb, Ids, CB, [], []),
+    Acc.
+
+get_ids(Docs) ->
+    lists:map(fun (#doc{id = Id}) -> Id end, Docs).
+
+
+create_mixed_docs(Db, Size) ->
+    fabric2_fdb:transactional(Db, fun (TxDb) ->
+        Docs = lists:map(fun (Id) ->
+            case Id rem 3 == 0 of
+                true -> create_local_doc(Id);
+                false -> create_doc(Id)
+            end
+        end, lists:seq(0, Size)),
+        {ok, _} = fabric2_db:update_docs(TxDb, Docs),
+        Docs
+    end).
+
+
+create_local_docs(Db, Size) ->
+    fabric2_fdb:transactional(Db, fun (TxDb) ->
+        Docs = lists:map(fun (Id) ->
+            create_local_doc(Id)
+        end, lists:seq(0, Size)),
+        {ok, _} = fabric2_db:update_docs(TxDb, Docs),
+        Docs
+    end).
+
+
+create_docs(Db, Size) ->
+    fabric2_fdb:transactional(Db, fun (TxDb) ->
+        Docs = lists:map(fun (Id) ->
+            create_doc(Id)
+        end, lists:seq(0, Size)),
+        {ok, _} = fabric2_db:update_docs(TxDb, Docs),
+        Docs
+    end).
+
+
+create_doc(Id) ->
+    couch_doc:from_json_obj({[
+        {<<"_id">>, list_to_binary([<<"doc-">>, integer_to_binary(Id)])},
+        {<<"value">>, 1}
+    ]}).
+
+
+create_local_doc(Id) ->
+    couch_doc:from_json_obj({[
+        {<<"_id">>, list_to_binary([<<"_local/doc-">>, integer_to_binary(Id)])},
+        {<<"value">>, 1}
+    ]}).
+
+
+shuffle(List) when is_list(List) ->
+    Tagged = [{rand:uniform(), Item} || Item <- List],
+    {_, Randomized} = lists:unzip(lists:sort(Tagged)),
+    Randomized.
diff --git a/test/elixir/test/all_docs_test.exs b/test/elixir/test/all_docs_test.exs
index d41d046b8..46ab1f8b3 100644
--- a/test/elixir/test/all_docs_test.exs
+++ b/test/elixir/test/all_docs_test.exs
@@ -276,6 +276,70 @@ defmodule AllDocsTest do
   end
 
   @tag :with_db
+  test "_local_docs POST with keys and limit", context do
+    expected = [
+      %{
+        "doc" => %{"_id" => "_local/one", "_rev" => "0-1", "value" => "one"},
+        "id" => "_local/one",
+        "key" => "_local/one",
+        "value" => %{"rev" => "0-1"}
+      },
+      %{
+        "doc" => %{"_id" => "_local/two", "_rev" => "0-1", "value" => "two"},
+        "id" => "_local/two",
+        "key" => "_local/two",
+        "value" => %{"rev" => "0-1"}
+      },
+      %{
+        "doc" => %{
+          "_id" => "three",
+          "_rev" => "1-878d3724976748bc881841046a276ceb",
+          "value" => "three"
+        },
+        "id" => "three",
+        "key" => "three",
+        "value" => %{"rev" => "1-878d3724976748bc881841046a276ceb"}
+      },
+      %{"error" => "not_found", "key" => "missing"},
+      %{"error" => "not_found", "key" => "_local/missing"}
+    ]
+
+    db_name = context[:db_name]
+
+    docs = [
+      %{
+        _id: "_local/one",
+        value: "one"
+      },
+      %{
+        _id: "_local/two",
+        value: "two"
+      },
+      %{
+        _id: "three",
+        value: "three"
+      }
+    ]
+
+    resp = Couch.post("/#{db_name}/_bulk_docs", body: %{docs: docs})
+    assert resp.status_code in [201, 202]
+
+    resp =
+      Couch.post(
+        "/#{db_name}/_all_docs",
+        body: %{
+          :keys => ["_local/one", "_local/two", "three", "missing", "_local/missing"],
+          :include_docs => true
+        }
+      )
+
+    assert resp.status_code == 200
+    rows = resp.body["rows"]
+    assert length(rows) == 5
+    assert rows == expected
+  end
+
+  @tag :with_db
   test "POST with query parameter and JSON body", context do
     db_name = context[:db_name]