Support partitioned queries

Co-authored-by: Garren Smith <garren.smith@gmail.com>
Co-authored-by: Peng Hui Jiang <jiangph@cn.ibm.com>

5 files changed, 166 insertions, 24 deletions

diff --git a/include/dreyfus.hrl b/include/dreyfus.hrl
index f86287c92..7c6a36945 100644
--- a/include/dreyfus.hrl
+++ b/include/dreyfus.hrl
@@ -32,6 +32,7 @@
 
 -record(index_query_args, {
     q,
+    partition=nil,
     limit=25,
     stale=false,
     include_docs=false,
diff --git a/src/dreyfus_httpd.erl b/src/dreyfus_httpd.erl
index 8db545466..e3ef5a96c 100644
--- a/src/dreyfus_httpd.erl
+++ b/src/dreyfus_httpd.erl
@@ -31,17 +31,10 @@ handle_search_req(#httpd{method=Method, path_parts=[_, _, _, _, IndexName]}=Req
     DbName = couch_db:name(Db),
     Start = os:timestamp(),
     QueryArgs = #index_query_args{
-        q = Query,
         include_docs = IncludeDocs,
         grouping = Grouping
     } = parse_index_params(Req),
-    case Query of
-        undefined ->
-            Msg = <<"Query must include a 'q' or 'query' argument">>,
-            throw({query_parse_error, Msg});
-        _ ->
-            ok
-    end,
+    validate_search_restrictions(Db, DDoc, QueryArgs),
     Response = case Grouping#grouping.by of
         nil ->
             case dreyfus_fabric_search:go(DbName, DDoc, IndexName, QueryArgs) of
@@ -206,6 +199,8 @@ parse_index_params(IndexParams) ->
 
 validate_index_query(q, Value, Args) ->
     Args#index_query_args{q=Value};
+validate_index_query(partition, Value, Args) ->
+    Args#index_query_args{partition=Value};
 validate_index_query(stale, Value, Args) ->
     Args#index_query_args{stale=Value};
 validate_index_query(limit, Value, Args) ->
@@ -254,6 +249,8 @@ parse_index_param("q", Value) ->
     [{q, ?l2b(Value)}];
 parse_index_param("query", Value) ->
     [{q, ?l2b(Value)}];
+parse_index_param("partition", Value) ->
+    [{partition, ?l2b(Value)}];
 parse_index_param("bookmark", Value) ->
     [{bookmark, ?l2b(Value)}];
 parse_index_param("sort", Value) ->
@@ -301,6 +298,8 @@ parse_json_index_param(<<"q">>, Value) ->
     [{q, Value}];
 parse_json_index_param(<<"query">>, Value) ->
     [{q, Value}];
+parse_json_index_param(<<"partition">>, Value) ->
+    [{partition, Value}];
 parse_json_index_param(<<"bookmark">>, Value) ->
     [{bookmark, Value}];
 parse_json_index_param(<<"sort">>, Value) ->
@@ -418,6 +417,63 @@ parse_non_negative_int_param(Name, Val, Prop, Default) ->
     end.
 
 
+validate_search_restrictions(Db, DDoc, Args) ->
+    #index_query_args{
+        q = Query,
+        partition = Partition,
+        grouping = Grouping
+    } = Args,
+    #grouping{
+        by = GroupBy
+    } = Grouping,
+
+    case Query of
+        undefined ->
+            Msg1 = <<"Query must include a 'q' or 'query' argument">>,
+            throw({query_parse_error, Msg1});
+        _ ->
+            ok
+    end,
+
+    DbPartitioned = fabric_util:is_partitioned(Db),
+    ViewPartitioned = get_view_partition_option(DDoc, DbPartitioned),
+
+    case not DbPartitioned andalso is_binary(Partition) of
+        true ->
+            Msg2 = <<"`partition` not supported on this index">>,
+            throw({bad_request, Msg2});
+        false ->
+            ok
+    end,
+
+    case {ViewPartitioned, is_binary(Partition)} of
+        {false, false} ->
+            ok;
+        {true, true} -> 
+            ok;
+        {true, false} ->
+            Msg3 = <<"`partition` parameter is mandatory "
+                        "for queries to this index.">>,
+            throw({bad_request, Msg3});
+        {false, true} ->
+            Msg4 = <<"`partition` not supported on this index">>,
+            throw({bad_request, Msg4})
+    end,
+
+    case GroupBy /= nil andalso is_binary(Partition) of
+        true ->
+            Msg5 = <<"`group_by` and `partition` are incompatible">>,
+            throw({bad_request, Msg5});
+        false ->
+            ok
+    end.
+
+
+get_view_partition_option(#doc{body = {Props}}, Default) ->
+    {Options} = couch_util:get_value(<<"options">>, Props, {[]}),
+    couch_util:get_value(<<"partitioned">>, Options, Default).
+
+
 hits_to_json(DbName, IncludeDocs, Hits) ->
     {Ids, HitData} = lists:unzip(lists:map(fun get_hit_data/1, Hits)),
     if IncludeDocs ->
diff --git a/src/dreyfus_index.erl b/src/dreyfus_index.erl
index 32ab40f08..e33a208ee 100644
--- a/src/dreyfus_index.erl
+++ b/src/dreyfus_index.erl
@@ -305,6 +305,7 @@ index_name(#index{dbname=DbName,ddoc_id=DDocId,name=IndexName}) ->
 args_to_proplist(#index_query_args{} = Args) ->
     [
      {'query', Args#index_query_args.q},
+     {partition, Args#index_query_args.partition},
      {limit, Args#index_query_args.limit},
      {refresh, Args#index_query_args.stale =:= false},
      {'after', Args#index_query_args.bookmark},
diff --git a/src/dreyfus_index_updater.erl b/src/dreyfus_index_updater.erl
index e2fbe2b05..40fd0c377 100644
--- a/src/dreyfus_index_updater.erl
+++ b/src/dreyfus_index_updater.erl
@@ -132,13 +132,19 @@ update_or_delete_index(IndexPid, Db, DI, Proc) ->
         true ->
             ok = clouseau_rpc:delete(IndexPid, Id);
         false ->
-            {ok, Doc} = couch_db:open_doc(Db, DI, []),
-            Json = couch_doc:to_json_obj(Doc, []),
-            [Fields|_] = proc_prompt(Proc, [<<"index_doc">>, Json]),
-            Fields1 = [list_to_tuple(Field) || Field <- Fields],
-            case Fields1 of
-                [] -> ok = clouseau_rpc:delete(IndexPid, Id);
-                _  -> ok = clouseau_rpc:update(IndexPid, Id, Fields1)
+            case maybe_skip_doc(Db, Id) of
+                true ->
+                    ok;
+                false ->
+                    {ok, Doc} = couch_db:open_doc(Db, DI, []),
+                    Json = couch_doc:to_json_obj(Doc, []),
+                    [Fields|_] = proc_prompt(Proc, [<<"index_doc">>, Json]),
+                    Fields1 = [list_to_tuple(Field) || Field <- Fields],
+                    Fields2 = maybe_add_partition(Db, Id, Fields1),
+                    case Fields2 of
+                        [] -> ok = clouseau_rpc:delete(IndexPid, Id);
+                        _  -> ok = clouseau_rpc:update(IndexPid, Id, Fields2)
+                    end
             end
     end.
 
@@ -157,3 +163,19 @@ update_task(NumChanges) ->
             (Changes2 * 100) div Total
     end,
     couch_task_status:update([{progress, Progress}, {changes_done, Changes2}]).
+
+maybe_skip_doc(Db, <<"_design/", _/binary>>) ->
+    couch_db:is_partitioned(Db);
+maybe_skip_doc(_Db, _Id) ->
+    false.
+
+maybe_add_partition(_Db, _Id, []) ->
+    [];
+maybe_add_partition(Db, Id, Fields) ->
+    case couch_db:is_partitioned(Db) of
+        true ->
+            Partition = couch_partition:from_docid(Id),
+            [{<<"_partition">>, Partition, {[]}} | Fields];
+        false ->
+            Fields
+    end.
diff --git a/src/dreyfus_util.erl b/src/dreyfus_util.erl
index 3b3f4f955..ae3133e7d 100644
--- a/src/dreyfus_util.erl
+++ b/src/dreyfus_util.erl
@@ -33,15 +33,31 @@
     verify_index_exists/2
 ]).
 
-get_shards(DbName, #index_query_args{stale=ok}) ->
-    mem3:ushards(DbName);
-get_shards(DbName, #index_query_args{stable=true}) ->
-    mem3:ushards(DbName);
-get_shards(DbName, #index_query_args{stale=false}) ->
-    mem3:shards(DbName);
+
+get_shards(DbName, #index_query_args{partition = nil} = Args) ->
+    case use_ushards(Args) of
+        true ->
+            mem3:ushards(DbName);
+        false ->
+            mem3:shards(DbName)
+    end;
+get_shards(DbName, #index_query_args{partition = Partition} = Args) ->
+    PartitionId = couch_partition:shard_key(Partition),
+    case use_ushards(Args) of
+        true ->
+            mem3:ushards(DbName, PartitionId);
+        false ->
+            mem3:shards(DbName, PartitionId)
+    end;
 get_shards(DbName, Args) ->
     get_shards(DbName, upgrade(Args)).
 
+use_ushards(#index_query_args{stale=ok}) ->
+    true;
+use_ushards(#index_query_args{stable=true}) ->
+    true;
+use_ushards(#index_query_args{}) ->
+    false.
 
 -spec sort(Order :: relevance | [any()], [#sortable{}]) -> [#sortable{}].
 sort(Sort, List0) ->
@@ -136,10 +152,34 @@ upgrade({index_query_args, Query, Limit, Stale, IncludeDocs, Bookmark,
         highlight_post_tag = HighlightPostTag,
         highlight_number = HighlightNumber,
         highlight_size = HighlightSize
+    };
+upgrade({index_query_args, Query, Limit, Stale, IncludeDocs, Bookmark,
+         Sort, Grouping, Stable, Counts, Ranges, Drilldown,
+         IncludeFields, HighlightFields, HighlightPreTag, HighlightPostTag,
+         HighlightNumber, HighlightSize, RawBookmark}) ->
+    #index_query_args{
+        q = Query,
+        limit = Limit,
+        stale = Stale,
+        include_docs = IncludeDocs,
+        bookmark = Bookmark,
+        sort = Sort,
+        grouping =  Grouping,
+        stable = Stable,
+        counts = Counts,
+        ranges = Ranges,
+        drilldown = Drilldown,
+        include_fields = IncludeFields,
+        highlight_fields = HighlightFields,
+        highlight_pre_tag = HighlightPreTag,
+        highlight_post_tag = HighlightPostTag,
+        highlight_number = HighlightNumber,
+        highlight_size = HighlightSize,
+        raw_bookmark = RawBookmark
     }.
 
-export(#index_query_args{counts = nil, ranges = nil, drilldown = [],
-    include_fields = nil, highlight_fields = nil} = Args) ->
+export(#index_query_args{partition = nil, counts = nil, ranges = nil,
+    drilldown = [], include_fields = nil, highlight_fields = nil} = Args) ->
     % Ensure existing searches work during the upgrade by creating an
     % #index_query_args record in the old format
     {index_query_args,
@@ -152,7 +192,8 @@ export(#index_query_args{counts = nil, ranges = nil, drilldown = [],
         Args#index_query_args.grouping,
         Args#index_query_args.stable
     };
-export(#index_query_args{include_fields = nil, highlight_fields = nil} = Args) ->
+export(#index_query_args{partition = nil, include_fields = nil,
+    highlight_fields = nil} = Args) ->
     {index_query_args,
         Args#index_query_args.q,
         Args#index_query_args.limit,
@@ -166,6 +207,27 @@ export(#index_query_args{include_fields = nil, highlight_fields = nil} = Args) -
         Args#index_query_args.ranges,
         Args#index_query_args.drilldown
     };
+export(#index_query_args{partition = nil} = Args) ->
+    {index_query_args,
+        Args#index_query_args.q,
+        Args#index_query_args.limit,
+        Args#index_query_args.stale,
+        Args#index_query_args.include_docs,
+        Args#index_query_args.bookmark,
+        Args#index_query_args.sort,
+        Args#index_query_args.grouping,
+        Args#index_query_args.stable,
+        Args#index_query_args.counts,
+        Args#index_query_args.ranges,
+        Args#index_query_args.drilldown,
+        Args#index_query_args.include_fields,
+        Args#index_query_args.highlight_fields,
+        Args#index_query_args.highlight_pre_tag,
+        Args#index_query_args.highlight_post_tag,
+        Args#index_query_args.highlight_number,
+        Args#index_query_args.highlight_size,
+        Args#index_query_args.raw_bookmark
+    };
 export(QueryArgs) ->
     QueryArgs.