Preserve user callbacks through DocumentOperationExecutorImpl

2 files changed, 16 insertions, 5 deletions

diff --git a/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/DocumentOperationExecutorImpl.java b/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/DocumentOperationExecutorImpl.java
index 98b58fc2b7c..824f6b901d3 100644
--- a/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/DocumentOperationExecutorImpl.java
+++ b/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/DocumentOperationExecutorImpl.java
@@ -166,22 +166,22 @@ public class DocumentOperationExecutorImpl implements DocumentOperationExecutor
 
     @Override
     public void get(DocumentId id, DocumentOperationParameters parameters, OperationContext context) {
-        accept(() -> asyncSession.get(id, parameters.withResponseHandler(handlerOf(context))), context);
+        accept(() -> asyncSession.get(id, parameters.withResponseHandler(handlerOf(parameters, context))), context);
     }
 
     @Override
     public void put(DocumentPut put, DocumentOperationParameters parameters, OperationContext context) {
-        accept(() -> asyncSession.put(put, parameters.withResponseHandler(handlerOf(context))), context);
+        accept(() -> asyncSession.put(put, parameters.withResponseHandler(handlerOf(parameters, context))), context);
     }
 
     @Override
     public void update(DocumentUpdate update, DocumentOperationParameters parameters, OperationContext context) {
-        accept(() -> asyncSession.update(update, parameters.withResponseHandler(handlerOf(context))), context);
+        accept(() -> asyncSession.update(update, parameters.withResponseHandler(handlerOf(parameters, context))), context);
     }
 
     @Override
     public void remove(DocumentId id, DocumentOperationParameters parameters, OperationContext context) {
-        accept(() -> asyncSession.remove(id, parameters.withResponseHandler(handlerOf(context))), context);
+        accept(() -> asyncSession.remove(id, parameters.withResponseHandler(handlerOf(parameters, context))), context);
     }
 
     @Override
@@ -230,8 +230,9 @@ public class DocumentOperationExecutorImpl implements DocumentOperationExecutor
         return resolveCluster(Optional.of(cluster), clusters).route();
     }
 
-    private ResponseHandler handlerOf(OperationContext context) {
+    private ResponseHandler handlerOf(DocumentOperationParameters parameters, OperationContext context) {
         return response -> {
+            parameters.responseHandler().ifPresent(originalHandler -> originalHandler.handleResponse(response));
             if (response.isSuccess())
                 context.success(response instanceof DocumentResponse ? Optional.ofNullable(((DocumentResponse) response).getDocument())
                                                                      : Optional.empty());
diff --git a/vespaclient-container-plugin/src/test/java/com/yahoo/document/restapi/DocumentOperationExecutorTest.java b/vespaclient-container-plugin/src/test/java/com/yahoo/document/restapi/DocumentOperationExecutorTest.java
index b57a0478ec2..331ad18e63c 100644
--- a/vespaclient-container-plugin/src/test/java/com/yahoo/document/restapi/DocumentOperationExecutorTest.java
+++ b/vespaclient-container-plugin/src/test/java/com/yahoo/document/restapi/DocumentOperationExecutorTest.java
@@ -232,6 +232,16 @@ public class DocumentOperationExecutorTest {
     }
 
     @Test
+    public void testCallback() {
+        AtomicBoolean called = new AtomicBoolean();
+        executor.get(doc1.getId(), parameters().withResponseHandler(__ -> called.set(true)), operationContext());
+        assertTrue(called.get());
+        assertEquals(List.of(), messages);
+        assertEquals(List.of(), errors);
+        assertEquals(List.of(), received);
+    }
+
+    @Test
     public void testVisit() throws InterruptedException {
         executor.put(new DocumentPut(doc1), parameters(), operationContext());
         executor.put(new DocumentPut(doc2), parameters(), operationContext());