Merge pull request #3573 from vespa-engine/vekterli/add-wanted-document-count-parameter-to-document-v1-visiting-api

Add wantedDocumentCount parameter to Document V1 visiting API

6 files changed, 146 insertions, 41 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 3e16bced996..e7ed9ce10db 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
@@ -8,11 +8,11 @@ import java.util.Optional;
 /**
  * Abstract the backend stuff for the REST API, such as retrieving or updating documents.
  *
- * @author dybis
+ * @author Haakon Dybdahl
  */
 public interface OperationHandler {
 
-    class VisitResult{
+    class VisitResult {
 
         public final Optional<String> token;
         public final String documentsAsJsonList;
@@ -23,7 +23,19 @@ public interface OperationHandler {
         }
     }
 
-    VisitResult visit(RestUri restUri, String documentSelection,  Optional<String> cluster, Optional<String> continuation) throws RestApiException;
+    class VisitOptions {
+        public final Optional<String> cluster;
+        public final Optional<String> continuation;
+        public final Optional<Integer> wantedDocumentCount;
+
+        public VisitOptions(Optional<String> cluster, Optional<String> continuation, Optional<Integer> wantedDocumentCount) {
+            this.cluster = cluster;
+            this.continuation = continuation;
+            this.wantedDocumentCount = wantedDocumentCount;
+        }
+    }
+
+    VisitResult visit(RestUri restUri, String documentSelection, VisitOptions options) throws RestApiException;
 
     void put(RestUri restUri, VespaXMLFeedReader.Operation data, Optional<String> route) throws RestApiException;
 
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 482a39c60e5..46678ea67e3 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
@@ -49,6 +49,7 @@ public class OperationHandlerImpl implements OperationHandler {
     }
 
     public static final int VISIT_TIMEOUT_MS = 120000;
+    public static final int WANTED_DOCUMENT_COUNT_UPPER_BOUND = 1000; // Approximates the max default size of a bucket
     private final DocumentAccess documentAccess;
     private final DocumentApiMetrics metricsHelper;
     private final ClusterEnumerator clusterEnumerator;
@@ -109,13 +110,8 @@ public class OperationHandlerImpl implements OperationHandler {
     }
 
     @Override
-    public VisitResult visit(
-            RestUri restUri,
-            String documentSelection,
-            Optional<String> cluster,
-            Optional<String> continuation) throws RestApiException {
-
-        VisitorParameters visitorParameters = createVisitorParameters(restUri, documentSelection, cluster, continuation);
+    public VisitResult visit(RestUri restUri, String documentSelection, VisitOptions options) throws RestApiException {
+        VisitorParameters visitorParameters = createVisitorParameters(restUri, documentSelection, options);
 
         VisitorControlHandler visitorControlHandler = new VisitorControlHandler();
         visitorParameters.setControlHandler(visitorControlHandler);
@@ -326,13 +322,13 @@ public class OperationHandlerImpl implements OperationHandler {
     private VisitorParameters createVisitorParameters(
             RestUri restUri,
             String documentSelection,
-            Optional<String> clusterName,
-            Optional<String> continuation)
+            VisitOptions options)
             throws RestApiException {
 
         StringBuilder selection = new StringBuilder();
 
         if (! documentSelection.isEmpty()) {
+            // TODO shouldn't selection be wrapped in () itself ?
             selection.append("(").append(documentSelection).append(" and ");
         }
         selection.append(restUri.getDocumentType()).append(" and (id.namespace=='").append(restUri.getNamespace()).append("')");
@@ -346,24 +342,26 @@ public class OperationHandlerImpl implements OperationHandler {
         params.setMaxBucketsPerVisitor(1);
         params.setMaxPending(32);
         params.setMaxFirstPassHits(1);
-        params.setMaxTotalHits(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.setToTimestamp(0L);
         params.setFromTimestamp(0L);
         params.setSessionTimeoutMs(VISIT_TIMEOUT_MS);
 
-        params.visitInconsistentBuckets(true);
+        params.visitInconsistentBuckets(true); // TODO document this as part of consistency doc
         params.setVisitorOrdering(VisitorOrdering.ASCENDING);
 
-        params.setRoute(resolveClusterRoute(clusterName));
+        params.setRoute(resolveClusterRoute(options.cluster));
 
         params.setTraceLevel(0);
         params.setPriority(DocumentProtocol.Priority.NORMAL_4);
         params.setVisitRemoves(false);
 
-        if (continuation.isPresent()) {
+        if (options.continuation.isPresent()) {
             try {
-                params.setResumeToken(ContinuationHit.getToken(continuation.get()));
+                params.setResumeToken(ContinuationHit.getToken(options.continuation.get()));
             } catch (Exception e) {
                 throw new RestApiException(Response.createErrorResponse(500, ExceptionUtils.getStackTrace(e), restUri, RestUri.apiErrorCodes.UNSPECIFIED));
             }
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 1f1eca9674b..5e0fea8ab7d 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
@@ -38,7 +38,7 @@ import java.util.concurrent.atomic.AtomicInteger;
 /**
  * API for handling single operation on a document and visiting.
  *
- * @author dybis
+ * @author Haakon Dybdahl
  */
 public class RestApi extends LoggingRequestHandler {
 
@@ -52,6 +52,7 @@ public class RestApi extends LoggingRequestHandler {
     private static final String SELECTION = "selection";
     private static final String CLUSTER = "cluster";
     private static final String CONTINUATION = "continuation";
+    private static final String WANTED_DOCUMENT_COUNT = "wantedDocumentCount";
     private static final String APPLICATION_JSON = "application/json";
     private final OperationHandler operationHandler;
     private SingleDocumentParser singleDocumentParser;
@@ -96,19 +97,33 @@ public class RestApi extends LoggingRequestHandler {
         this.singleDocumentParser = new SingleDocumentParser(docTypeManager);
     }
 
-    // Returns null if invalid value.
-    private Optional<Boolean> parseBoolean(String parameter, HttpRequest request) {
+    private static Optional<String> requestProperty(String parameter, HttpRequest request) {
         final String property = request.getProperty(parameter);
         if (property != null && ! property.isEmpty()) {
-            switch (property) {
-                case "true" : return Optional.of(true);
-                case "false": return Optional.of(false);
-                default : return null;
-            }
+            return Optional.of(property);
         }
         return Optional.empty();
     }
 
+    private static boolean parseBooleanStrict(String value) {
+        if ("true".equalsIgnoreCase(value)) {
+            return true;
+        } else if ("false".equalsIgnoreCase(value)) {
+            return false;
+        }
+        throw new IllegalArgumentException(String.format("Value not convertible to bool: '%s'", value));
+    }
+
+    private static Optional<Boolean> parseBoolean(String parameter, HttpRequest request) {
+        Optional<String> property = requestProperty(parameter, request);
+        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);
+    }
+
     @Override
     public HttpResponse handle(HttpRequest request) {
         try {
@@ -134,8 +149,10 @@ public class RestApi extends LoggingRequestHandler {
             return Response.createErrorResponse(500, "Exception while parsing URI: " + e2.getMessage(), RestUri.apiErrorCodes.URL_PARSING);
         }
 
-        Optional<Boolean> create = parseBoolean(CREATE_PARAMETER_NAME, request);
-        if (create == null) {
+        final Optional<Boolean> create;
+        try {
+            create = parseBoolean(CREATE_PARAMETER_NAME, request);
+        } catch (IllegalArgumentException e) {
             return Response.createErrorResponse(403, "Non valid value for 'create' parameter, must be empty, true, or " +
                     "false: " + request.getProperty(CREATE_PARAMETER_NAME), RestUri.apiErrorCodes.INVALID_CREATE_VALUE);
         }
@@ -184,9 +201,7 @@ public class RestApi extends LoggingRequestHandler {
         if (condition != null && ! condition.isEmpty()) {
             operationUpdate.getDocumentUpdate().setCondition(new TestAndSetCondition(condition));
         }
-        if (create.isPresent()) {
-            operationUpdate.getDocumentUpdate().setCreateIfNonExistent(create.get());
-        }
+        create.ifPresent(c -> operationUpdate.getDocumentUpdate().setCreateIfNonExistent(c));
         return operationUpdate;
     }
 
@@ -214,11 +229,16 @@ 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);
+    }
     
     private HttpResponse handleVisit(RestUri restUri, HttpRequest request) throws RestApiException {
         String documentSelection = Optional.ofNullable(request.getProperty(SELECTION)).orElse("");
         if (restUri.getGroup().isPresent() && ! restUri.getGroup().get().value.isEmpty()) {
             if (! documentSelection.isEmpty()) {
+                // TODO why is this restriction in place? Document selection allows composition of location predicate and other expressions
                 return Response.createErrorResponse(
                         400,
                         "Visiting does not support setting value for group/value in combination with expression, try using only expression parameter instead.",
@@ -234,11 +254,17 @@ public class RestApi extends LoggingRequestHandler {
         }
         Optional<String> cluster = Optional.ofNullable(request.getProperty(CLUSTER));
         Optional<String> continuation = Optional.ofNullable(request.getProperty(CONTINUATION));
-        final OperationHandler.VisitResult visit = operationHandler.visit(restUri, documentSelection, cluster, continuation);
-        final ObjectNode resultNode = mapper.createObjectNode();
-        if (visit.token.isPresent()) {
-            resultNode.put(CONTINUATION, visit.token.get());
+        Optional<Integer> wantedDocumentCount;
+        try {
+            wantedDocumentCount = parseInteger(WANTED_DOCUMENT_COUNT, request);
+        } catch (IllegalArgumentException e) {
+            return createInvalidParameterResponse(WANTED_DOCUMENT_COUNT, "Expected integer");
         }
+
+        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));
         resultNode.putArray(DOCUMENTS).addPOJO(visit.documentsAsJsonList);
         resultNode.put(PATH_NAME, restUri.getRawPath());
 
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 796c6d23deb..5735e84f3fe 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
@@ -91,7 +91,7 @@ public class OperationHandlerImplTest {
         return new String( stream.toByteArray());
     }
 
-    private class OperationHandlerImplFixture {
+    private static class OperationHandlerImplFixture {
         DocumentAccess documentAccess = mock(DocumentAccess.class);
         AtomicReference<VisitorParameters> assignedParameters = new AtomicReference<>();
         VisitorControlHandler.CompletionCode completionCode = VisitorControlHandler.CompletionCode.SUCCESS;
@@ -123,6 +123,14 @@ public class OperationHandlerImplTest {
         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));
+    }
+
+    private static OperationHandler.VisitOptions emptyVisitOptions() {
+        return new OperationHandler.VisitOptions(Optional.empty(), Optional.empty(), Optional.empty());
+    }
+
     @Test
     public void timeout_without_buckets_visited_throws_timeout_error() throws Exception {
         OperationHandlerImplFixture fixture = new OperationHandlerImplFixture();
@@ -131,7 +139,7 @@ public class OperationHandlerImplTest {
         // RestApiException hides its guts internally, so cannot trivially use @Rule directly to check for error category
         try {
             OperationHandlerImpl handler = fixture.createHandler();
-            handler.visit(dummyVisitUri(), "", Optional.empty(), Optional.empty());
+            handler.visit(dummyVisitUri(), "", emptyVisitOptions());
         } catch (RestApiException e) {
             assertThat(e.getResponse().getStatus(), is(500));
             assertThat(renderRestApiExceptionAsString(e), containsString("Timed out"));
@@ -145,7 +153,7 @@ public class OperationHandlerImplTest {
         fixture.bucketsVisited = 1;
 
         OperationHandlerImpl handler = fixture.createHandler();
-        handler.visit(dummyVisitUri(), "", Optional.empty(), Optional.empty());
+        handler.visit(dummyVisitUri(), "", emptyVisitOptions());
     }
 
     @Test
@@ -153,8 +161,50 @@ public class OperationHandlerImplTest {
         OperationHandlerImplFixture fixture = new OperationHandlerImplFixture();
         OperationHandlerImpl handler = fixture.createHandler();
 
-        handler.visit(dummyVisitUri(), "", Optional.empty(), Optional.empty());
+        handler.visit(dummyVisitUri(), "", emptyVisitOptions());
 
         assertThat(fixture.assignedParameters.get().getSessionTimeoutMs(), is((long)OperationHandlerImpl.VISIT_TIMEOUT_MS));
     }
+
+    private static VisitorParameters generatedParametersFromVisitOptions(OperationHandler.VisitOptions options) throws Exception {
+        OperationHandlerImplFixture fixture = new OperationHandlerImplFixture();
+        OperationHandlerImpl handler = fixture.createHandler();
+
+        handler.visit(dummyVisitUri(), "", options);
+        return fixture.assignedParameters.get();
+    }
+
+    @Test
+    public void provided_wanted_document_count_is_propagated_to_visitor_parameters() throws Exception {
+        VisitorParameters params = generatedParametersFromVisitOptions(visitOptionsWithWantedDocumentCount(123));
+        assertThat(params.getMaxTotalHits(), is((long)123));
+    }
+
+    @Test
+    public void wanted_document_count_is_1_unless_specified() throws Exception {
+        VisitorParameters params = generatedParametersFromVisitOptions(emptyVisitOptions());
+        assertThat(params.getMaxTotalHits(), is((long)1));
+    }
+
+    @Test
+    public void too_low_wanted_document_count_is_bounded_to_1() throws Exception {
+        VisitorParameters params = generatedParametersFromVisitOptions(visitOptionsWithWantedDocumentCount(-1));
+        assertThat(params.getMaxTotalHits(), is((long)1));
+
+        params = generatedParametersFromVisitOptions(visitOptionsWithWantedDocumentCount(Integer.MIN_VALUE));
+        assertThat(params.getMaxTotalHits(), is((long)1));
+
+        params = generatedParametersFromVisitOptions(visitOptionsWithWantedDocumentCount(0));
+        assertThat(params.getMaxTotalHits(), is((long)1));
+    }
+
+    @Test
+    public void too_high_wanted_document_count_is_bounded_to_upper_bound() throws Exception {
+        VisitorParameters params = generatedParametersFromVisitOptions(visitOptionsWithWantedDocumentCount(OperationHandlerImpl.WANTED_DOCUMENT_COUNT_UPPER_BOUND + 1));
+        assertThat(params.getMaxTotalHits(), is((long)OperationHandlerImpl.WANTED_DOCUMENT_COUNT_UPPER_BOUND));
+
+        params = generatedParametersFromVisitOptions(visitOptionsWithWantedDocumentCount(Integer.MAX_VALUE));
+        assertThat(params.getMaxTotalHits(), is((long)OperationHandlerImpl.WANTED_DOCUMENT_COUNT_UPPER_BOUND));
+    }
+
 }
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 97f45c4062a..f353013232f 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
@@ -18,9 +18,11 @@ public class MockedOperationHandler implements OperationHandler {
     int deleteCount = 0;
 
     @Override
-    public VisitResult visit(RestUri restUri, String documentSelection, Optional<String> cluster, Optional<String> continuation) throws RestApiException {
-        return new VisitResult(Optional.of("token"), "List of json docs, cont token " + continuation.map(a->a).orElse("not set") + ", doc selection: '"
-                + documentSelection + "'");
+    public VisitResult visit(RestUri restUri, String documentSelection, VisitOptions options) throws RestApiException {
+        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(""));
     }
 
     @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 10ae80a5d03..91390e3a0d8 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
@@ -270,6 +270,7 @@ public class RestApiTest {
         assertThat(rest, containsString(visit_response_part3));
     }
 
+    // TODO why is this a limitation?
     String visit_test_bad_uri = "/document/v1/namespace/document-type/group/abc?continuation=abc&selection=foo";
     String visit_test_bad_response = "Visiting does not support setting value for group/value in combination with expression";
 
@@ -294,6 +295,22 @@ public class RestApiTest {
         assertThat(rest, containsString(visit_test_response_selection_rewrite));
     }
 
+    @Test
+    public void wanted_document_count_returned_parameter_is_propagated() throws IOException {
+        Request request = new Request(String.format("http://localhost:%s/document/v1/namespace/document-type/docid/?wantedDocumentCount=321", getFirstListenPort()));
+        HttpGet get = new HttpGet(request.getUri());
+        String rest = doRest(get);
+        assertThat(rest, containsString("min docs returned: 321"));
+    }
+
+    @Test
+    public void 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"));
+    }
+
     private String doRest(HttpRequestBase request) throws IOException {
         HttpClient client = HttpClientBuilder.create().build();
         HttpResponse response = client.execute(request);