Merge pull request #4873 from vespa-engine/vekterli/add-fieldset-and-concurrency-support-to-document-v1-visiting-api

Add fieldset and concurrency support to document v1 visiting api

6 files changed, 178 insertions, 19 deletions

diff --git a/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/OperationHandler.java b/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/OperationHandler.java
index e7ed9ce10db..6a811eb5afd 100644
--- a/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/OperationHandler.java
+++ b/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/OperationHandler.java
@@ -27,11 +27,66 @@ public interface OperationHandler {
         public final Optional<String> cluster;
         public final Optional<String> continuation;
         public final Optional<Integer> wantedDocumentCount;
+        public final Optional<String> fieldSet;
+        public final Optional<Integer> concurrency;
 
+        /** @deprecated Use a VisitOptions.Builder instead */
+        @Deprecated
         public VisitOptions(Optional<String> cluster, Optional<String> continuation, Optional<Integer> wantedDocumentCount) {
             this.cluster = cluster;
             this.continuation = continuation;
             this.wantedDocumentCount = wantedDocumentCount;
+            this.fieldSet = Optional.empty();
+            this.concurrency = Optional.empty();
+        }
+
+        private VisitOptions(Builder builder) {
+            this.cluster = Optional.ofNullable(builder.cluster);
+            this.continuation = Optional.ofNullable(builder.continuation);
+            this.wantedDocumentCount = Optional.ofNullable(builder.wantedDocumentCount);
+            this.fieldSet = Optional.ofNullable(builder.fieldSet);
+            this.concurrency = Optional.ofNullable(builder.concurrency);
+        }
+
+        public static class Builder {
+            String cluster;
+            String continuation;
+            Integer wantedDocumentCount;
+            String fieldSet;
+            Integer concurrency;
+
+            public Builder cluster(String cluster) {
+                this.cluster = cluster;
+                return this;
+            }
+
+            public Builder continuation(String continuation) {
+                this.continuation = continuation;
+                return this;
+            }
+
+            public Builder wantedDocumentCount(Integer count) {
+                this.wantedDocumentCount = count;
+                return this;
+            }
+
+            public Builder fieldSet(String fieldSet) {
+                this.fieldSet = fieldSet;
+                return this;
+            }
+
+            public Builder concurrency(Integer concurrency) {
+                this.concurrency = concurrency;
+                return this;
+            }
+
+            public VisitOptions build() {
+                return new VisitOptions(this);
+            }
+        }
+
+        public static Builder builder() {
+            return new Builder();
         }
     }
 
diff --git a/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/OperationHandlerImpl.java b/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/OperationHandlerImpl.java
index 8fe2007e88a..d6ae4a9285a 100644
--- a/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/OperationHandlerImpl.java
+++ b/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/OperationHandlerImpl.java
@@ -385,14 +385,14 @@ public class OperationHandlerImpl implements OperationHandler {
 
         VisitorParameters params = new VisitorParameters(selection.toString());
         // Only return fieldset that is part of the document.
-        params.fieldSet(restUri.getDocumentType() + ":[document]");
+        params.fieldSet(options.fieldSet.orElse(restUri.getDocumentType() + ":[document]"));
         params.setMaxBucketsPerVisitor(1);
         params.setMaxPending(32);
         params.setMaxFirstPassHits(1);
         params.setMaxTotalHits(options.wantedDocumentCount
                 .map(n -> Math.min(Math.max(n, 1), WANTED_DOCUMENT_COUNT_UPPER_BOUND))
                 .orElse(1));
-        params.setThrottlePolicy(new StaticThrottlePolicy().setMaxPendingCount(1));
+        params.setThrottlePolicy(new StaticThrottlePolicy().setMaxPendingCount(options.concurrency.orElse(1)));
         params.setToTimestamp(0L);
         params.setFromTimestamp(0L);
         params.setSessionTimeoutMs(VISIT_TIMEOUT_MS);
diff --git a/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/resource/RestApi.java b/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/resource/RestApi.java
index d5d6a0dc40b..a45557859ad 100644
--- a/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/resource/RestApi.java
+++ b/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/resource/RestApi.java
@@ -6,7 +6,6 @@ import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.inject.Inject;
 
 import com.fasterxml.jackson.databind.node.ObjectNode;
-import com.yahoo.jdisc.Metric;
 import com.yahoo.container.handler.ThreadpoolConfig;
 import com.yahoo.container.jdisc.HttpRequest;
 import com.yahoo.container.jdisc.HttpResponse;
@@ -54,6 +53,8 @@ public class RestApi extends LoggingRequestHandler {
     private static final String CLUSTER = "cluster";
     private static final String CONTINUATION = "continuation";
     private static final String WANTED_DOCUMENT_COUNT = "wantedDocumentCount";
+    private static final String FIELD_SET = "fieldSet";
+    private static final String CONCURRENCY = "concurrency";
     private static final String APPLICATION_JSON = "application/json";
     private final OperationHandler operationHandler;
     private SingleDocumentParser singleDocumentParser;
@@ -122,9 +123,12 @@ public class RestApi extends LoggingRequestHandler {
         return property.map(RestApi::parseBooleanStrict);
     }
 
-    private static Optional<Integer> parseInteger(String parameter, HttpRequest request) throws NumberFormatException {
-        Optional<String> property = requestProperty(parameter, request);
-        return property.map(Integer::parseInt);
+    private static int parsePositiveInt(String str) throws NumberFormatException {
+        int parsed = Integer.parseInt(str);
+        if (parsed <= 0) {
+            throw new IllegalArgumentException("Parsed number was negative or zero");
+        }
+        return parsed;
     }
 
     @Override
@@ -236,6 +240,43 @@ public class RestApi extends LoggingRequestHandler {
     private static HttpResponse createInvalidParameterResponse(String parameter, String explanation) {
         return Response.createErrorResponse(403, String.format("Invalid '%s' value. %s", parameter, explanation), RestUri.apiErrorCodes.UNSPECIFIED);
     }
+
+    static class BadRequestParameterException extends IllegalArgumentException {
+        private String parameter;
+
+        BadRequestParameterException(String parameter, String message) {
+            super(message);
+            this.parameter = parameter;
+        }
+
+        String getParameter() {
+            return parameter;
+        }
+    }
+
+    private static Optional<Integer> parsePositiveIntegerRequestParameter(String parameter, HttpRequest request) {
+        Optional<String> property = requestProperty(parameter, request);
+        if (!property.isPresent()) {
+            return Optional.empty();
+        }
+        try {
+            return property.map(RestApi::parsePositiveInt);
+        } catch (IllegalArgumentException e) {
+            throw new BadRequestParameterException(parameter, "Expected positive integer");
+        }
+    }
+
+    private static OperationHandler.VisitOptions visitOptionsFromRequest(HttpRequest request) {
+        final OperationHandler.VisitOptions.Builder optionsBuilder = OperationHandler.VisitOptions.builder();
+
+        Optional.ofNullable(request.getProperty(CLUSTER)).ifPresent(c -> optionsBuilder.cluster(c));
+        Optional.ofNullable(request.getProperty(CONTINUATION)).ifPresent(c -> optionsBuilder.continuation(c));
+        Optional.ofNullable(request.getProperty(FIELD_SET)).ifPresent(fs -> optionsBuilder.fieldSet(fs));
+        parsePositiveIntegerRequestParameter(WANTED_DOCUMENT_COUNT, request).ifPresent(c -> optionsBuilder.wantedDocumentCount(c));
+        parsePositiveIntegerRequestParameter(CONCURRENCY, request).ifPresent(c -> optionsBuilder.concurrency(c));
+
+        return optionsBuilder.build();
+    }
     
     private HttpResponse handleVisit(RestUri restUri, HttpRequest request) throws RestApiException {
         String documentSelection = Optional.ofNullable(request.getProperty(SELECTION)).orElse("");
@@ -255,16 +296,12 @@ public class RestApi extends LoggingRequestHandler {
                 documentSelection = "id.group='" + group.value + "'";
             }
         }
-        Optional<String> cluster = Optional.ofNullable(request.getProperty(CLUSTER));
-        Optional<String> continuation = Optional.ofNullable(request.getProperty(CONTINUATION));
-        Optional<Integer> wantedDocumentCount;
+        OperationHandler.VisitOptions options;
         try {
-            wantedDocumentCount = parseInteger(WANTED_DOCUMENT_COUNT, request);
-        } catch (IllegalArgumentException e) {
-            return createInvalidParameterResponse(WANTED_DOCUMENT_COUNT, "Expected integer");
+            options = visitOptionsFromRequest(request);
+        } catch (BadRequestParameterException e) {
+            return createInvalidParameterResponse(e.getParameter(), e.getMessage());
         }
-
-        final OperationHandler.VisitOptions options = new OperationHandler.VisitOptions(cluster, continuation, wantedDocumentCount);
         final OperationHandler.VisitResult visit = operationHandler.visit(restUri, documentSelection, options);
         final ObjectNode resultNode = mapper.createObjectNode();
         visit.token.ifPresent(t -> resultNode.put(CONTINUATION, t));
diff --git a/vespaclient-container-plugin/src/test/java/com/yahoo/document/restapi/OperationHandlerImplTest.java b/vespaclient-container-plugin/src/test/java/com/yahoo/document/restapi/OperationHandlerImplTest.java
index 61bbe1e7030..3a6782f5830 100644
--- a/vespaclient-container-plugin/src/test/java/com/yahoo/document/restapi/OperationHandlerImplTest.java
+++ b/vespaclient-container-plugin/src/test/java/com/yahoo/document/restapi/OperationHandlerImplTest.java
@@ -6,6 +6,7 @@ import com.yahoo.documentapi.ProgressToken;
 import com.yahoo.documentapi.VisitorControlHandler;
 import com.yahoo.documentapi.VisitorParameters;
 import com.yahoo.documentapi.VisitorSession;
+import com.yahoo.messagebus.StaticThrottlePolicy;
 import com.yahoo.metrics.simple.MetricReceiver;
 import com.yahoo.vdslib.VisitorStatistics;
 import com.yahoo.vespaclient.ClusterDef;
@@ -18,7 +19,9 @@ import java.util.*;
 import java.util.concurrent.atomic.AtomicReference;
 
 import static org.hamcrest.CoreMatchers.containsString;
+import static org.hamcrest.CoreMatchers.instanceOf;
 import static org.hamcrest.core.Is.is;
+import static org.hamcrest.core.IsEqual.equalTo;
 import static org.junit.Assert.*;
 import static org.mockito.Matchers.any;
 import static org.mockito.Mockito.mock;
@@ -127,16 +130,20 @@ public class OperationHandlerImplTest {
         }
     }
 
+    private static OperationHandler.VisitOptions.Builder optionsBuilder() {
+        return OperationHandler.VisitOptions.builder();
+    }
+
     private static RestUri dummyVisitUri() throws Exception {
         return new RestUri(new URI("http://localhost/document/v1/namespace/document-type/docid/"));
     }
 
     private static OperationHandler.VisitOptions visitOptionsWithWantedDocumentCount(int wantedDocumentCount) {
-        return new OperationHandler.VisitOptions(Optional.empty(), Optional.empty(), Optional.of(wantedDocumentCount));
+        return optionsBuilder().wantedDocumentCount(wantedDocumentCount).build();
     }
 
     private static OperationHandler.VisitOptions emptyVisitOptions() {
-        return new OperationHandler.VisitOptions(Optional.empty(), Optional.empty(), Optional.empty());
+        return optionsBuilder().build();
     }
 
     @Test
@@ -242,4 +249,30 @@ public class OperationHandlerImplTest {
         assertThat(params.getMaxTotalHits(), is((long)OperationHandlerImpl.WANTED_DOCUMENT_COUNT_UPPER_BOUND));
     }
 
+    @Test
+    public void field_set_covers_all_fields_by_default() throws Exception {
+        VisitorParameters params = generatedParametersFromVisitOptions(emptyVisitOptions());
+        assertThat(params.fieldSet(), equalTo("document-type:[document]"));
+    }
+
+    @Test
+    public void provided_fieldset_is_propagated_to_visitor_parameters() throws Exception {
+        VisitorParameters params = generatedParametersFromVisitOptions(optionsBuilder().fieldSet("document-type:bjarne").build());
+        assertThat(params.fieldSet(), equalTo("document-type:bjarne"));
+    }
+
+    @Test
+    public void concurrency_is_1_by_default() throws Exception {
+        VisitorParameters params = generatedParametersFromVisitOptions(emptyVisitOptions());
+        assertThat(params.getThrottlePolicy(), instanceOf(StaticThrottlePolicy.class));
+        assertThat(((StaticThrottlePolicy)params.getThrottlePolicy()).getMaxPendingCount(), is((int)1));
+    }
+
+    @Test
+    public void concurrency_is_propagated_to_visitor_parameters() throws Exception {
+        VisitorParameters params = generatedParametersFromVisitOptions(optionsBuilder().concurrency(3).build());
+        assertThat(params.getThrottlePolicy(), instanceOf(StaticThrottlePolicy.class));
+        assertThat(((StaticThrottlePolicy)params.getThrottlePolicy()).getMaxPendingCount(), is((int)3));
+    }
+
 }
diff --git a/vespaclient-container-plugin/src/test/java/com/yahoo/document/restapi/resource/MockedOperationHandler.java b/vespaclient-container-plugin/src/test/java/com/yahoo/document/restapi/resource/MockedOperationHandler.java
index f353013232f..3b16d9a378c 100644
--- a/vespaclient-container-plugin/src/test/java/com/yahoo/document/restapi/resource/MockedOperationHandler.java
+++ b/vespaclient-container-plugin/src/test/java/com/yahoo/document/restapi/resource/MockedOperationHandler.java
@@ -22,7 +22,9 @@ public class MockedOperationHandler implements OperationHandler {
         return new VisitResult(Optional.of("token"), "List of json docs, cont token "
                 + options.continuation.orElse("not set") + ", doc selection: '"
                 + documentSelection + "'"
-                + options.wantedDocumentCount.map(n -> String.format(", min docs returned: %d", n)).orElse(""));
+                + options.wantedDocumentCount.map(n -> String.format(", min docs returned: %d", n)).orElse("")
+                + options.fieldSet.map(s -> String.format(", field set: '%s'", s)).orElse("")
+                + options.concurrency.map(n -> String.format(", concurrency: %d", n)).orElse(""));
     }
 
     @Override
diff --git a/vespaclient-container-plugin/src/test/java/com/yahoo/document/restapi/resource/RestApiTest.java b/vespaclient-container-plugin/src/test/java/com/yahoo/document/restapi/resource/RestApiTest.java
index 91390e3a0d8..faee5947794 100644
--- a/vespaclient-container-plugin/src/test/java/com/yahoo/document/restapi/resource/RestApiTest.java
+++ b/vespaclient-container-plugin/src/test/java/com/yahoo/document/restapi/resource/RestApiTest.java
@@ -304,11 +304,43 @@ public class RestApiTest {
     }
 
     @Test
-    public void wanted_document_count_parameter_returns_error_response() throws IOException {
+    public void invalid_wanted_document_count_parameter_returns_error_response() throws IOException {
         Request request = new Request(String.format("http://localhost:%s/document/v1/namespace/document-type/docid/?wantedDocumentCount=aardvark", getFirstListenPort()));
         HttpGet get = new HttpGet(request.getUri());
         String rest = doRest(get);
-        assertThat(rest, containsString("Invalid 'wantedDocumentCount' value. Expected integer"));
+        assertThat(rest, containsString("Invalid 'wantedDocumentCount' value. Expected positive integer"));
+    }
+
+    @Test
+    public void negative_document_count_parameter_returns_error_response() throws IOException {
+        Request request = new Request(String.format("http://localhost:%s/document/v1/namespace/document-type/docid/?wantedDocumentCount=-1", getFirstListenPort()));
+        HttpGet get = new HttpGet(request.getUri());
+        String rest = doRest(get);
+        assertThat(rest, containsString("Invalid 'wantedDocumentCount' value. Expected positive integer"));
+    }
+
+    @Test
+    public void fieldset_parameter_is_propagated() throws IOException {
+        Request request = new Request(String.format("http://localhost:%s/document/v1/namespace/document-type/docid/?fieldSet=foo,baz", getFirstListenPort()));
+        HttpGet get = new HttpGet(request.getUri());
+        String rest = doRest(get);
+        assertThat(rest, containsString("field set: 'foo,baz'"));
+    }
+
+    @Test
+    public void concurrency_parameter_is_propagated() throws IOException {
+        Request request = new Request(String.format("http://localhost:%s/document/v1/namespace/document-type/docid/?concurrency=42", getFirstListenPort()));
+        HttpGet get = new HttpGet(request.getUri());
+        String rest = doRest(get);
+        assertThat(rest, containsString("concurrency: 42"));
+    }
+
+    @Test
+    public void invalid_concurrency_parameter_returns_error_response() throws IOException {
+        Request request = new Request(String.format("http://localhost:%s/document/v1/namespace/document-type/docid/?concurrency=badgers", getFirstListenPort()));
+        HttpGet get = new HttpGet(request.getUri());
+        String rest = doRest(get);
+        assertThat(rest, containsString("Invalid 'concurrency' value. Expected positive integer"));
     }
 
     private String doRest(HttpRequestBase request) throws IOException {