Add route option, to be tested in system test.

4 files changed, 30 insertions, 12 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 fb1b317399c..1a9c8c46693 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
@@ -25,11 +25,11 @@ public interface OperationHandler {
 
     VisitResult visit(RestUri restUri, String documentSelection,  Optional<String> cluster, Optional<String> continuation) throws RestApiException;
 
-    void put(RestUri restUri, VespaXMLFeedReader.Operation data) throws RestApiException;
+    void put(RestUri restUri, VespaXMLFeedReader.Operation data, Optional<String> route) throws RestApiException;
 
-    void update(RestUri restUri, VespaXMLFeedReader.Operation data) throws RestApiException;
+    void update(RestUri restUri, VespaXMLFeedReader.Operation data, Optional<String> route) throws RestApiException;
 
-    void delete(RestUri restUri, String condition) throws RestApiException;
+    void delete(RestUri restUri, String condition, Optional<String> route) throws RestApiException;
 
     Optional<String> get(RestUri restUri) 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 5ad7524ca2f..93b6dfcab71 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
@@ -14,6 +14,7 @@ import com.yahoo.documentapi.SyncSession;
 import com.yahoo.documentapi.VisitorControlHandler;
 import com.yahoo.documentapi.VisitorParameters;
 import com.yahoo.documentapi.VisitorSession;
+import com.yahoo.documentapi.messagebus.MessageBusSyncSession;
 import com.yahoo.documentapi.messagebus.protocol.DocumentProtocol;
 import com.yahoo.messagebus.StaticThrottlePolicy;
 import com.yahoo.storage.searcher.ContinuationHit;
@@ -140,12 +141,24 @@ public class OperationHandlerImpl implements OperationHandler {
         throw new RestApiException(Response.createErrorResponse(500, localDataVisitorHandler.getErrors(), restUri));
     }
 
+    private void setRoute(SyncSession session, Optional<String> route) throws RestApiException {
+        if (route.isPresent()) {
+            if (! (session instanceof MessageBusSyncSession)) {
+                // Not sure if this ever could happen but better be safe.
+                throw new RestApiException(Response.createErrorResponse(
+                        400, "Can not set route since the API is not using message bus."));
+            }
+            ((MessageBusSyncSession) session).setRoute(route.get());
+        }
+    }
+
     @Override
-    public void put(RestUri restUri, VespaXMLFeedReader.Operation data) throws RestApiException {
+    public void put(RestUri restUri, VespaXMLFeedReader.Operation data, Optional<String> route) throws RestApiException {
         SyncSession syncSession = syncSessions.alloc();
         try {
             DocumentPut put = new DocumentPut(data.getDocument());
             put.setCondition(data.getCondition());
+            setRoute(syncSession, route);
             syncSession.put(put);
         } catch (DocumentAccessException documentException) {
             throw new RestApiException(createErrorResponse(documentException, restUri));
@@ -157,8 +170,9 @@ public class OperationHandlerImpl implements OperationHandler {
     }
 
     @Override
-    public void update(RestUri restUri, VespaXMLFeedReader.Operation data) throws RestApiException {
+    public void update(RestUri restUri, VespaXMLFeedReader.Operation data, Optional<String> route) throws RestApiException {
         SyncSession syncSession = syncSessions.alloc();
+        setRoute(syncSession, route);
         try {
             syncSession.update(data.getDocumentUpdate());
         } catch (DocumentAccessException documentException) {
@@ -171,8 +185,9 @@ public class OperationHandlerImpl implements OperationHandler {
     }
 
     @Override
-    public void delete(RestUri restUri, String condition) throws RestApiException {
+    public void delete(RestUri restUri, String condition, Optional<String> route) throws RestApiException {
         SyncSession syncSession = syncSessions.alloc();
+        setRoute(syncSession, route);
         try {
             DocumentId id = new DocumentId(restUri.generateFullId());
             DocumentRemove documentRemove = new DocumentRemove(id);
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 80769909b9f..1cffe9d8ff6 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
@@ -43,6 +43,7 @@ public class RestApi extends LoggingRequestHandler {
 
     private static final String CREATE_PARAMETER_NAME = "create";
     private static final String CONDITION_PARAMETER_NAME = "condition";
+    private static final String ROUTE_PARAMETER_NAME = "route";
     private static final String DOCUMENTS = "documents";
     private static final String FIELDS = "fields";
     private static final String DOC_ID_NAME = "id";
@@ -138,19 +139,21 @@ public class RestApi extends LoggingRequestHandler {
                     "false: " + request.getProperty(CREATE_PARAMETER_NAME));
         }
         String condition = request.getProperty(CONDITION_PARAMETER_NAME);
+        Optional<String> route = Optional.ofNullable(request.getProperty(ROUTE_PARAMETER_NAME));
+
         Optional<ObjectNode> resultJson = Optional.empty();
         try {
             switch (request.getMethod()) {
                 case GET:    // Vespa Visit/Get
                     return restUri.getDocId().isEmpty() ? handleVisit(restUri, request) : handleGet(restUri);
                 case POST:   // Vespa Put
-                    operationHandler.put(restUri, createPutOperation(request, restUri.generateFullId(), condition));
+                    operationHandler.put(restUri, createPutOperation(request, restUri.generateFullId(), condition), route);
                     break;
                 case PUT:    // Vespa Update
-                    operationHandler.update(restUri, createUpdateOperation(request, restUri.generateFullId(), condition, create));
+                    operationHandler.update(restUri, createUpdateOperation(request, restUri.generateFullId(), condition, create), route);
                     break;
                 case DELETE: // Vespa Delete
-                    operationHandler.delete(restUri, condition);
+                    operationHandler.delete(restUri, condition, route);
                     break;
                 default:
                     return new Response(405, Optional.empty(), Optional.of(restUri));
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 ab5b36c74d1..2d98b72e579 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
@@ -24,13 +24,13 @@ public class MockedOperationHandler implements OperationHandler {
     }
 
     @Override
-    public void put(RestUri restUri, VespaXMLFeedReader.Operation data) throws RestApiException {
+    public void put(RestUri restUri, VespaXMLFeedReader.Operation data, Optional<String> route) throws RestApiException {
         log.append("PUT: " + data.getDocument().getId());
         log.append(data.getDocument().getBody().toString());
     }
 
     @Override
-    public void update(RestUri restUri, VespaXMLFeedReader.Operation data) throws RestApiException {
+    public void update(RestUri restUri, VespaXMLFeedReader.Operation data, Optional<String> route) throws RestApiException {
         log.append("UPDATE: " + data.getDocumentUpdate().getId());
         log.append(data.getDocumentUpdate().getFieldUpdates().toString());
         if (data.getDocumentUpdate().getCreateIfNonExistent()) {
@@ -39,7 +39,7 @@ public class MockedOperationHandler implements OperationHandler {
     }
 
     @Override
-    public void delete(RestUri restUri, String condition) throws RestApiException {
+    public void delete(RestUri restUri, String condition, Optional<String> route) throws RestApiException {
         deleteCount++;
         if (deleteCount == 2) {
             String theLog = log.toString();