Include document count in responses

2 files changed, 13 insertions, 1 deletions

diff --git a/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/resource/DocumentV1ApiHandler.java b/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/resource/DocumentV1ApiHandler.java
index 7c80db2ddc3..7b2578b4094 100644
--- a/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/resource/DocumentV1ApiHandler.java
+++ b/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/resource/DocumentV1ApiHandler.java
@@ -645,6 +645,10 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
             json.writeStringField("message", message);
         }
 
+        synchronized void writeDocumentCount(long count) throws IOException {
+            json.writeNumberField("documentCount", count);
+        }
+
         synchronized void writeDocId(DocumentId id) throws IOException {
             json.writeStringField("id", id.toString());
         }
@@ -1083,11 +1087,17 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
                                     if (getProgress() != null && ! getProgress().isFinished())
                                         response.writeContinuation(getProgress().serializeToString());
 
+                                    if (getVisitorStatistics() != null)
+                                        response.writeDocumentCount(getVisitorStatistics().getDocumentsReturned());
+
                                     response.respond(Response.Status.OK);
                                     break;
                                 }
                             default:
                                 response.writeMessage(error.get() != null ? error.get() : message != null ? message : "Visiting failed");
+                                if (getVisitorStatistics() != null)
+                                    response.writeDocumentCount(getVisitorStatistics().getDocumentsReturned());
+
                                 response.respond(Response.Status.INTERNAL_SERVER_ERROR);
                         }
                         visitDispatcher.execute(() -> {
diff --git a/vespaclient-container-plugin/src/test/java/com/yahoo/document/restapi/resource/DocumentV1ApiTest.java b/vespaclient-container-plugin/src/test/java/com/yahoo/document/restapi/resource/DocumentV1ApiTest.java
index cee02618cd7..c297c09f906 100644
--- a/vespaclient-container-plugin/src/test/java/com/yahoo/document/restapi/resource/DocumentV1ApiTest.java
+++ b/vespaclient-container-plugin/src/test/java/com/yahoo/document/restapi/resource/DocumentV1ApiTest.java
@@ -211,6 +211,7 @@ public class DocumentV1ApiTest {
             parameters.getLocalDataHandler().onMessage(new PutDocumentMessage(new DocumentPut(doc3)), tokens.get(2));
             VisitorStatistics statistics = new VisitorStatistics();
             statistics.setBucketsVisited(1);
+            statistics.setDocumentsReturned(3);
             parameters.getControlHandler().onVisitorStatistics(statistics);
             parameters.getControlHandler().onDone(VisitorControlHandler.CompletionCode.TIMEOUT, "timeout is OK");
         });
@@ -235,7 +236,8 @@ public class DocumentV1ApiTest {
                        "     \"id\": \"id:space:music:g=a:three\"," +
                        "     \"fields\": {}" +
                        "    }" +
-                       "  ]" +
+                       "  ]," +
+                       "  \"documentCount\": 3" +
                        "}", response.readAll());
         assertEquals(200, response.getStatus());