Merge pull request #4557 from apache/nouveau-partition-support

finish partitioned support for nouveau

6 files changed, 106 insertions, 7 deletions

diff --git a/src/nouveau/src/nouveau_api.erl b/src/nouveau/src/nouveau_api.erl
index 5bf6b1731..7744eb161 100644
--- a/src/nouveau/src/nouveau_api.erl
+++ b/src/nouveau/src/nouveau_api.erl
@@ -24,7 +24,7 @@
     delete_path/1,
     delete_path/2,
     delete_doc/3,
-    update_doc/4,
+    update_doc/5,
     search/2
 ]).
 
@@ -110,10 +110,17 @@ delete_doc(#index{} = Index, DocId, UpdateSeq) when
             send_error(Reason)
     end.
 
-update_doc(#index{} = Index, DocId, UpdateSeq, Fields) when
-    is_binary(DocId), is_integer(UpdateSeq), is_list(Fields)
+update_doc(#index{} = Index, DocId, UpdateSeq, Partition, Fields) when
+    is_binary(DocId),
+    is_integer(UpdateSeq),
+    (is_binary(Partition) orelse Partition == null),
+    is_list(Fields)
 ->
-    ReqBody = {[{<<"seq">>, UpdateSeq}, {<<"fields">>, Fields}]},
+    ReqBody = #{
+        seq => UpdateSeq,
+        partition => Partition,
+        fields => Fields
+    },
     Resp = send_if_enabled(
         doc_url(Index, DocId), [?JSON_CONTENT_TYPE], put, jiffy:encode(ReqBody)
     ),
diff --git a/src/nouveau/src/nouveau_fabric_search.erl b/src/nouveau/src/nouveau_fabric_search.erl
index ca101ba9c..5212bf620 100644
--- a/src/nouveau/src/nouveau_fabric_search.erl
+++ b/src/nouveau/src/nouveau_fabric_search.erl
@@ -36,7 +36,7 @@ go(DbName, GroupId, IndexName, QueryArgs0) when is_binary(GroupId) ->
     go(DbName, DDoc, IndexName, QueryArgs0);
 go(DbName, #doc{} = DDoc, IndexName, QueryArgs0) ->
     {ok, Index} = nouveau_util:design_doc_to_index(DbName, DDoc, IndexName),
-    Shards = mem3:shards(DbName),
+    Shards = get_shards(DbName, QueryArgs0),
     {PackedBookmark, #{limit := Limit, sort := Sort} = QueryArgs1} =
         maps:take(bookmark, QueryArgs0),
     Bookmark = nouveau_bookmark:unpack(DbName, PackedBookmark),
@@ -219,3 +219,9 @@ merge_facets(null, FacetsB, _Limit) ->
 merge_facets(FacetsA, FacetsB, _Limit) ->
     Combiner = fun(_, V1, V2) -> nouveau_maps:merge_with(fun(_, V3, V4) -> V3 + V4 end, V1, V2) end,
     nouveau_maps:merge_with(Combiner, FacetsA, FacetsB).
+
+get_shards(DbName, #{partition := Partition}) when is_binary(Partition) ->
+    PartitionId = couch_partition:shard_key(Partition),
+    mem3:shards(DbName, PartitionId);
+get_shards(DbName, _QueryArgs) ->
+    mem3:shards(DbName).
diff --git a/src/nouveau/src/nouveau_httpd.erl b/src/nouveau/src/nouveau_httpd.erl
index 8d27048a1..5cb50de6a 100644
--- a/src/nouveau/src/nouveau_httpd.erl
+++ b/src/nouveau/src/nouveau_httpd.erl
@@ -65,6 +65,7 @@ handle_search_req_int(#httpd{method = 'GET', path_parts = [_, _, _, _, IndexName
     DbName = couch_db:name(Db),
     QueryArgs = validate_query_args(#{
         query => chttpd:qs_value(Req, "q"),
+        partition => chttpd:qs_value(Req, "partition"),
         limit => chttpd:qs_value(Req, "limit"),
         sort => chttpd:qs_value(Req, "sort"),
         ranges => chttpd:qs_value(Req, "ranges"),
@@ -82,6 +83,7 @@ handle_search_req_int(
     ReqBody = chttpd:json_body(Req, [return_maps]),
     QueryArgs = validate_query_args(#{
         query => maps:get(<<"q">>, ReqBody, undefined),
+        partition => chttpd:qs_value(Req, "partition"),
         limit => maps:get(<<"limit">>, ReqBody, undefined),
         sort => json_or_undefined(<<"sort">>, ReqBody),
         ranges => json_or_undefined(<<"ranges">>, ReqBody),
@@ -175,6 +177,10 @@ validate_query_arg(query, undefined) ->
     throw({query_parse_error, <<"q parameter is mandatory">>});
 validate_query_arg(query, Val) when is_list(Val); is_binary(Val) ->
     couch_util:to_binary(Val);
+validate_query_arg(partition, undefined) ->
+    null;
+validate_query_arg(partition, Val) when is_list(Val); is_binary(Val) ->
+    couch_util:to_binary(Val);
 validate_query_arg(limit, undefined) ->
     25;
 validate_query_arg(limit, Limit) when is_integer(Limit), Limit > 0 ->
diff --git a/src/nouveau/src/nouveau_index_updater.erl b/src/nouveau/src/nouveau_index_updater.erl
index af39faecf..a78b1ff15 100644
--- a/src/nouveau/src/nouveau_index_updater.erl
+++ b/src/nouveau/src/nouveau_index_updater.erl
@@ -86,11 +86,18 @@ load_docs(FDI, {Db, Index, Proc, ChangesDone, TotalChanges}) ->
             {ok, Doc} = couch_db:open_doc(Db, DI, []),
             Json = couch_doc:to_json_obj(Doc, []),
             [Fields | _] = proc_prompt(Proc, [<<"nouveau_index_doc">>, Json]),
+            Partition =
+                case couch_db:is_partitioned(Db) of
+                    true ->
+                        couch_partition:from_docid(Id);
+                    false ->
+                        null
+                end,
             case Fields of
                 [] ->
                     ok = nouveau_api:delete_doc(Index, Id, Seq);
                 _ ->
-                    case nouveau_api:update_doc(Index, Id, Seq, Fields) of
+                    case nouveau_api:update_doc(Index, Id, Seq, Partition, Fields) of
                         ok ->
                             ok;
                         {error, Reason} ->
diff --git a/test/elixir/test/config/nouveau.elixir b/test/elixir/test/config/nouveau.elixir
index fdcc66de2..5c13aac2b 100644
--- a/test/elixir/test/config/nouveau.elixir
+++ b/test/elixir/test/config/nouveau.elixir
@@ -14,6 +14,9 @@
     "ranges",
     "ranges (open)",
     "mango search by number",
-    "mango search by string"
+    "mango search by string",
+    "search GET (partitioned)",
+    "search POST (partitioned)",
+    "mango (partitioned)"
   ]
 }
diff --git a/test/elixir/test/nouveau_test.exs b/test/elixir/test/nouveau_test.exs
index 33e3f66db..3bea874d9 100644
--- a/test/elixir/test/nouveau_test.exs
+++ b/test/elixir/test/nouveau_test.exs
@@ -20,6 +20,19 @@ defmodule NouveauTest do
     assert resp.status_code in [201]
   end
 
+  def create_partitioned_search_docs(db_name) do
+    resp = Couch.post("/#{db_name}/_bulk_docs",
+      headers: ["Content-Type": "application/json"],
+      body: %{:docs => [
+                %{"_id" => "foo:doc4", "foo" => "foo", "bar" => 42},
+                %{"_id" => "bar:doc3", "foo" => "bar", "bar" => 12.0},
+                %{"_id" => "foo:doc1", "foo" => "baz", "bar" => 0},
+                %{"_id" => "bar:doc2", "foo" => "foobar", "bar" => 100},
+      ]}
+    )
+    assert resp.status_code in [201]
+  end
+
   def create_ddoc(db_name, opts \\ %{}) do
     default_ddoc = %{
       nouveau: %{
@@ -290,4 +303,61 @@ defmodule NouveauTest do
     assert ids == ["doc4"]
   end
 
+  @tag :with_partitioned_db
+  test "search GET (partitioned)", context do
+    db_name = context[:db_name]
+    create_partitioned_search_docs(db_name)
+    create_ddoc(db_name)
+
+    url = "/#{db_name}/_partition/foo/_design/foo/_nouveau/bar"
+    resp = Couch.get(url, query: %{q: "*:*", include_docs: true})
+    assert_status_code(resp, 200)
+    ids = get_ids(resp)
+    assert ids == ["foo:doc1", "foo:doc4"]
+
+    url = "/#{db_name}/_partition/bar/_design/foo/_nouveau/bar"
+    resp = Couch.get(url, query: %{q: "*:*", include_docs: true})
+    assert_status_code(resp, 200)
+    ids = get_ids(resp)
+    assert ids == ["bar:doc2", "bar:doc3"]
+  end
+
+  @tag :with_partitioned_db
+  test "search POST (partitioned)", context do
+    db_name = context[:db_name]
+    create_partitioned_search_docs(db_name)
+    create_ddoc(db_name)
+
+    url = "/#{db_name}/_partition/foo/_design/foo/_nouveau/bar"
+    resp = Couch.post(url, body: %{q: "*:*", include_docs: true})
+    assert_status_code(resp, 200)
+    ids = get_ids(resp)
+    assert ids == ["foo:doc1", "foo:doc4"]
+
+    url = "/#{db_name}/_partition/bar/_design/foo/_nouveau/bar"
+    resp = Couch.post(url, body: %{q: "*:*", include_docs: true})
+    assert_status_code(resp, 200)
+    ids = get_ids(resp)
+    assert ids == ["bar:doc2", "bar:doc3"]
+  end
+
+  @tag :with_partitioned_db
+  test "mango (partitioned)", context do
+    db_name = context[:db_name]
+    create_partitioned_search_docs(db_name)
+    create_mango_index(db_name)
+
+    url = "/#{db_name}/_partition/foo/_find"
+    resp = Couch.post(url, body: %{selector: %{foo: %{"$eq": "foo"}}})
+    assert_status_code(resp, 200)
+    ids = get_mango_ids(resp)
+    assert ids == ["foo:doc4"]
+
+    url = "/#{db_name}/_partition/bar/_find"
+    resp = Couch.post(url, body: %{selector: %{foo: %{"$eq": "bar"}}})
+    assert_status_code(resp, 200)
+    ids = get_mango_ids(resp)
+    assert ids == ["bar:doc3"]
+  end
+
 end