Add wantedDocumentCount parameter to Document V1 visiting API

Visiting will continue until the provided number of documents has
been returned, or the session times out.
Parameter is bounded by an implementation defined maximum.
Remove old visit() interface method, as it should not have been covered
by an ExportPackage and therefore no one should have managed to
actually use it externally.

This fixes #3394

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..9134805da98 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.map(a->a).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);