Allow specifying the return type on Update

src/main/java/com/arangodb/ArangoCollection.java

@@ -271,6 +271,23 @@ <T> MultiDocumentEntity<DocumentUpdateEntity<T>> replaceDocuments(
     <T> DocumentUpdateEntity<T> updateDocument(String key, T value, DocumentUpdateOptions options)
             throws ArangoDBException;
 
+    /**
+     * Partially updates the document identified by document-key. The value must contain a document with the attributes
+     * to patch (the patch document). All attributes from the patch document will be added to the existing document if
+     * they do not yet exist, and overwritten in the existing document if they do exist there.
+     *
+     * @param key           The key of the document
+     * @param value         A representation of a single document (POJO, VPackSlice or String for JSON)
+     * @param options       Additional options, can be null
+     * @param returnType    Type of the returned newDocument and/or oldDocument
+     * @return information about the document
+     * @throws ArangoDBException
+     * @see <a href="https://www.arangodb.com/docs/stable/http/document-working-with-documents.html#update-document">API
+     * Documentation</a>
+     */
+    <T, U> DocumentUpdateEntity<U> updateDocument(String key, T value, DocumentUpdateOptions options, Class<U> returnType)
+            throws ArangoDBException;
+
     /**
      * Partially updates documents, the documents to update are specified by the _key attributes in the objects on
      * values. Vales must contain a list of document updates with the attributes to patch (the patch documents). All
@@ -301,6 +318,23 @@ <T> DocumentUpdateEntity<T> updateDocument(String key, T value, DocumentUpdateOp
     <T> MultiDocumentEntity<DocumentUpdateEntity<T>> updateDocuments(
             Collection<T> values, DocumentUpdateOptions options) throws ArangoDBException;
 
+    /**
+     * Partially updates documents, the documents to update are specified by the _key attributes in the objects on
+     * values. Vales must contain a list of document updates with the attributes to patch (the patch documents). All
+     * attributes from the patch documents will be added to the existing documents if they do not yet exist, and
+     * overwritten in the existing documents if they do exist there.
+     *
+     * @param values     A list of documents (POJO, VPackSlice or String for JSON)
+     * @param options    Additional options, can be null
+     * @param returnType Type of the returned newDocument and/or oldDocument
+     * @return information about the documents
+     * @throws ArangoDBException
+     * @see <a href="https://www.arangodb.com/docs/stable/http/document-working-with-documents.html#update-documents">API
+     * Documentation</a>
+     */
+    <T, U> MultiDocumentEntity<DocumentUpdateEntity<U>> updateDocuments(
+            Collection<T> values, DocumentUpdateOptions options, Class<U> returnType) throws ArangoDBException;
+
     /**
      * Deletes the document with the given {@code key} from the collection.
      *

src/main/java/com/arangodb/async/ArangoCollectionAsync.java

@@ -265,6 +265,25 @@ <T> CompletableFuture<DocumentUpdateEntity<T>> updateDocument(
             final T value,
             final DocumentUpdateOptions options);
 
+    /**
+     * Partially updates the document identified by document-key. The value must contain a document with the attributes
+     * to patch (the patch document). All attributes from the patch document will be added to the existing document if
+     * they do not yet exist, and overwritten in the existing document if they do exist there.
+     *
+     * @param key           The key of the document
+     * @param value         A representation of a single document (POJO, VPackSlice or String for Json)
+     * @param options       Additional options, can be null
+     * @param returnType    Type of the returned newDocument and/or oldDocument
+     * @return information about the document
+     * @see <a href="https://www.arangodb.com/docs/stable/http/document-working-with-documents.html#update-document">API
+     * Documentation</a>
+     */
+    <T, U> CompletableFuture<DocumentUpdateEntity<U>> updateDocument(
+            final String key,
+            final T value,
+            final DocumentUpdateOptions options,
+            final Class<U> returnType);
+
     /**
      * Partially updates documents, the documents to update are specified by the _key attributes in the objects on
      * values. Vales must contain a list of document updates with the attributes to patch (the patch documents). All
@@ -294,6 +313,24 @@ <T> CompletableFuture<MultiDocumentEntity<DocumentUpdateEntity<T>>> updateDocume
             final Collection<T> values,
             final DocumentUpdateOptions options);
 
+    /**
+     * Partially updates documents, the documents to update are specified by the _key attributes in the objects on
+     * values. Vales must contain a list of document updates with the attributes to patch (the patch documents). All
+     * attributes from the patch documents will be added to the existing documents if they do not yet exist, and
+     * overwritten in the existing documents if they do exist there.
+     *
+     * @param values     A list of documents (POJO, VPackSlice or String for Json)
+     * @param options    Additional options, can be null
+     * @param returnType Type of the returned newDocument and/or oldDocument
+     * @return information about the documents
+     * @see <a href="https://www.arangodb.com/docs/stable/http/document-working-with-documents.html#update-documents">API
+     * Documentation</a>
+     */
+    <T, U> CompletableFuture<MultiDocumentEntity<DocumentUpdateEntity<U>>> updateDocuments(
+            final Collection<T> values,
+            final DocumentUpdateOptions options,
+            final Class<U> returnType);
+
     /**
      * Removes a document
      *

src/main/java/com/arangodb/async/internal/ArangoCollectionAsyncImpl.java

@@ -167,35 +167,48 @@ public <T> CompletableFuture<MultiDocumentEntity<DocumentUpdateEntity<T>>> repla
 
     @Override
     public <T> CompletableFuture<DocumentUpdateEntity<T>> updateDocument(final String key, final T value) {
-        final DocumentUpdateOptions options = new DocumentUpdateOptions();
-        return executor.execute(updateDocumentRequest(key, value, options),
-                updateDocumentResponseDeserializer(value, options));
+        return updateDocument(key, value, new DocumentUpdateOptions());
     }
 
     @Override
     public <T> CompletableFuture<DocumentUpdateEntity<T>> updateDocument(
             final String key,
             final T value,
             final DocumentUpdateOptions options) {
+        return updateDocument(key, value, options, (Class<T>) value.getClass());
+    }
+
+    @Override
+    public <T, U> CompletableFuture<DocumentUpdateEntity<U>> updateDocument(
+            final String key,
+            final T value,
+            final DocumentUpdateOptions options,
+            final Class<U> returnType) {
         return executor.execute(updateDocumentRequest(key, value, options),
-                updateDocumentResponseDeserializer(value, options));
+                updateDocumentResponseDeserializer(value, options, returnType));
     }
 
     @Override
     public <T> CompletableFuture<MultiDocumentEntity<DocumentUpdateEntity<T>>> updateDocuments(
             final Collection<T> values) {
-        final DocumentUpdateOptions params = new DocumentUpdateOptions();
-        return executor.execute(updateDocumentsRequest(values, params),
-                updateDocumentsResponseDeserializer(values, params));
+        return updateDocuments(values, new DocumentUpdateOptions());
     }
 
     @Override
     public <T> CompletableFuture<MultiDocumentEntity<DocumentUpdateEntity<T>>> updateDocuments(
             final Collection<T> values,
             final DocumentUpdateOptions options) {
+        return updateDocuments(values, options, values.isEmpty() ? null : (Class<T>) values.iterator().next().getClass());
+    }
+
+    @Override
+    public <T, U> CompletableFuture<MultiDocumentEntity<DocumentUpdateEntity<U>>> updateDocuments(
+            final Collection<T> values,
+            final DocumentUpdateOptions options,
+            final Class<U> returnType) {
         final DocumentUpdateOptions params = (options != null ? options : new DocumentUpdateOptions());
         return executor.execute(updateDocumentsRequest(values, params),
-                updateDocumentsResponseDeserializer(values, params));
+                updateDocumentsResponseDeserializer(returnType));
     }
 
     @Override

src/main/java/com/arangodb/internal/ArangoCollectionImpl.java

@@ -168,8 +168,14 @@ public <T> DocumentUpdateEntity<T> updateDocument(final String key, final T valu
     @Override
     public <T> DocumentUpdateEntity<T> updateDocument(
             final String key, final T value, final DocumentUpdateOptions options) throws ArangoDBException {
+        return updateDocument(key, value, options, (Class<T>) value.getClass());
+    }
+
+    @Override
+    public <T, U> DocumentUpdateEntity<U> updateDocument(
+            final String key, final T value, final DocumentUpdateOptions options, final Class<U> returnType) throws ArangoDBException {
         return executor.execute(updateDocumentRequest(key, value, options),
-                updateDocumentResponseDeserializer(value, options));
+                updateDocumentResponseDeserializer(value, options, returnType));
     }
 
     @Override
@@ -181,9 +187,15 @@ public <T> MultiDocumentEntity<DocumentUpdateEntity<T>> updateDocuments(final Co
     @Override
     public <T> MultiDocumentEntity<DocumentUpdateEntity<T>> updateDocuments(
             final Collection<T> values, final DocumentUpdateOptions options) throws ArangoDBException {
+        return updateDocuments(values, options, values.isEmpty() ? null : (Class<T>) values.iterator().next().getClass());
+    }
+
+    @Override
+    public <T, U> MultiDocumentEntity<DocumentUpdateEntity<U>> updateDocuments(
+            final Collection<T> values, final DocumentUpdateOptions options, final Class<U> returnType) throws ArangoDBException {
         final DocumentUpdateOptions params = (options != null ? options : new DocumentUpdateOptions());
         return executor
-                .execute(updateDocumentsRequest(values, params), updateDocumentsResponseDeserializer(values, params));
+                .execute(updateDocumentsRequest(values, params), updateDocumentsResponseDeserializer(returnType));
     }
 
     @Override

src/main/java/com/arangodb/internal/InternalArangoCollection.java

@@ -357,18 +357,18 @@ protected <T> Request updateDocumentRequest(final String key, final T value, fin
         return request;
     }
 
-    protected <T> ResponseDeserializer<DocumentUpdateEntity<T>> updateDocumentResponseDeserializer(
-            final T value, final DocumentUpdateOptions options) {
+    protected <T, U> ResponseDeserializer<DocumentUpdateEntity<U>> updateDocumentResponseDeserializer(
+            final T value, final DocumentUpdateOptions options, final Class<U> returnType) {
         return response -> {
             final VPackSlice body = response.getBody();
-            final DocumentUpdateEntity<T> doc = util().deserialize(body, DocumentUpdateEntity.class);
+            final DocumentUpdateEntity<U> doc = util().deserialize(body, DocumentUpdateEntity.class);
             final VPackSlice newDoc = body.get(NEW);
             if (newDoc.isObject()) {
-                doc.setNew(util(Serializer.CUSTOM).deserialize(newDoc, value.getClass()));
+                doc.setNew(util(Serializer.CUSTOM).deserialize(newDoc, returnType));
             }
             final VPackSlice oldDoc = body.get(OLD);
             if (oldDoc.isObject()) {
-                doc.setOld(util(Serializer.CUSTOM).deserialize(oldDoc, value.getClass()));
+                doc.setOld(util(Serializer.CUSTOM).deserialize(oldDoc, returnType));
             }
             if (options == null || Boolean.TRUE != options.getSilent()) {
                 final Map<DocumentField.Type, String> values = new HashMap<>();
@@ -399,14 +399,8 @@ protected <T> Request updateDocumentsRequest(final Collection<T> values, final D
 
     @SuppressWarnings("unchecked")
     protected <T> ResponseDeserializer<MultiDocumentEntity<DocumentUpdateEntity<T>>> updateDocumentsResponseDeserializer(
-            final Collection<T> values, final DocumentUpdateOptions params) {
+            final Class<T> returnType) {
         return response -> {
-            Class<T> type = null;
-            if (Boolean.TRUE == params.getReturnNew() || Boolean.TRUE == params.getReturnOld()) {
-                if (!values.isEmpty()) {
-                    type = (Class<T>) values.iterator().next().getClass();
-                }
-            }
             final MultiDocumentEntity<DocumentUpdateEntity<T>> multiDocument = new MultiDocumentEntity<>();
             final Collection<DocumentUpdateEntity<T>> docs = new ArrayList<>();
             final Collection<ErrorEntity> errors = new ArrayList<>();
@@ -423,11 +417,11 @@ protected <T> ResponseDeserializer<MultiDocumentEntity<DocumentUpdateEntity<T>>>
                         final DocumentUpdateEntity<T> doc = util().deserialize(next, DocumentUpdateEntity.class);
                         final VPackSlice newDoc = next.get(NEW);
                         if (newDoc.isObject()) {
-                            doc.setNew(util(Serializer.CUSTOM).deserialize(newDoc, type));
+                            doc.setNew(util(Serializer.CUSTOM).deserialize(newDoc, returnType));
                         }
                         final VPackSlice oldDoc = next.get(OLD);
                         if (oldDoc.isObject()) {
-                            doc.setOld(util(Serializer.CUSTOM).deserialize(oldDoc, type));
+                            doc.setOld(util(Serializer.CUSTOM).deserialize(oldDoc, returnType));
                         }
                         docs.add(doc);
                         documentsAndErrors.add(doc);

src/test/java/com/arangodb/ArangoCollectionTest.java

@@ -435,6 +435,23 @@ public void updateDocument() {
         assertThat(readResult.getProperties().keySet(), hasItem("c"));
     }
 
+    @Test
+    public void updateDocumentWithDifferentReturnType() {
+        final String key = "key-" + UUID.randomUUID().toString();
+        final BaseDocument doc = new BaseDocument(key);
+        doc.addAttribute("a", "test");
+        collection.insertDocument(doc);
+
+        final DocumentUpdateEntity<BaseDocument> updateResult = collection
+                .updateDocument(key, Collections.singletonMap("b", "test"), new DocumentUpdateOptions().returnNew(true), BaseDocument.class);
+        assertThat(updateResult, is(notNullValue()));
+        assertThat(updateResult.getKey(), is(key));
+        BaseDocument updated = updateResult.getNew();
+        assertThat(updated, is(notNullValue()));
+        assertThat(updated.getAttribute("a"), is("test"));
+        assertThat(updated.getAttribute("b"), is("test"));
+    }
+
     @Test
     public void updateDocumentUpdateRev() {
         final BaseDocument doc = new BaseDocument();
@@ -2194,6 +2211,36 @@ public void updateDocuments() {
         assertThat(updateResult.getErrors().size(), is(0));
     }
 
+    @Test
+    public void updateDocumentsWithDifferentReturnType() {
+        List<String> keys = IntStream.range(0, 3).mapToObj(it -> "key-" + UUID.randomUUID().toString()).collect(Collectors.toList());
+        List<BaseDocument> docs = keys.stream()
+                .map(BaseDocument::new)
+                .peek(it -> it.addAttribute("a", "test"))
+                .collect(Collectors.toList());
+
+        collection.insertDocuments(docs);
+
+        List<Map<String, Object>> modifiedDocs = docs.stream()
+                .peek(it -> it.addAttribute("b", "test"))
+                .map(it -> {
+                    Map<String, Object> map = new HashMap<>();
+                    map.put("_key", it.getKey());
+                    map.put("a", it.getAttribute("a"));
+                    map.put("b", it.getAttribute("b"));
+                    return map;
+                })
+                .collect(Collectors.toList());
+
+        final MultiDocumentEntity<DocumentUpdateEntity<BaseDocument>> updateResult = collection
+                .updateDocuments(modifiedDocs, new DocumentUpdateOptions().returnNew(true), BaseDocument.class);
+        assertThat(updateResult.getDocuments().size(), is(3));
+        assertThat(updateResult.getErrors().size(), is(0));
+        assertThat(updateResult.getDocuments().stream().map(DocumentUpdateEntity::getNew)
+                        .allMatch(it -> it.getAttribute("a").equals("test") && it.getAttribute("b").equals("test")),
+                is(true));
+    }
+
     @Test
     public void updateDocumentsOne() {
         final Collection<BaseDocument> values = new ArrayList<>();

src/test/java/com/arangodb/async/ArangoCollectionTest.java

@@ -31,10 +31,12 @@
 import java.util.*;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 
 import static org.hamcrest.CoreMatchers.notNullValue;
-import static org.hamcrest.Matchers.*;
 import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.*;
 import static org.junit.Assert.fail;
 import static org.junit.Assume.assumeTrue;
 
@@ -284,6 +286,24 @@ public void updateDocument() throws InterruptedException, ExecutionException {
         assertThat(readResult.getProperties().keySet(), hasItem("c"));
     }
 
+    @Test
+    public void updateDocumentWithDifferentReturnType() throws ExecutionException, InterruptedException {
+        ArangoCollectionAsync collection = db.collection(COLLECTION_NAME);
+        final String key = "key-" + UUID.randomUUID().toString();
+        final BaseDocument doc = new BaseDocument(key);
+        doc.addAttribute("a", "test");
+        collection.insertDocument(doc).get();
+
+        final DocumentUpdateEntity<BaseDocument> updateResult = collection
+                .updateDocument(key, Collections.singletonMap("b", "test"), new DocumentUpdateOptions().returnNew(true), BaseDocument.class).get();
+        assertThat(updateResult, is(notNullValue()));
+        assertThat(updateResult.getKey(), is(key));
+        BaseDocument updated = updateResult.getNew();
+        assertThat(updated, is(notNullValue()));
+        assertThat(updated.getAttribute("a"), is("test"));
+        assertThat(updated.getAttribute("b"), is("test"));
+    }
+
     @Test
     public void updateDocumentIfMatch() throws InterruptedException, ExecutionException {
         final BaseDocument doc = new BaseDocument();
@@ -1842,6 +1862,37 @@ public void updateDocuments() throws InterruptedException, ExecutionException {
                 .get();
     }
 
+    @Test
+    public void updateDocumentsWithDifferentReturnType() throws ExecutionException, InterruptedException {
+        ArangoCollectionAsync collection = db.collection(COLLECTION_NAME);
+        List<String> keys = IntStream.range(0, 3).mapToObj(it -> "key-" + UUID.randomUUID().toString()).collect(Collectors.toList());
+        List<BaseDocument> docs = keys.stream()
+                .map(BaseDocument::new)
+                .peek(it -> it.addAttribute("a", "test"))
+                .collect(Collectors.toList());
+
+        collection.insertDocuments(docs).get();
+
+        List<Map<String, Object>> modifiedDocs = docs.stream()
+                .peek(it -> it.addAttribute("b", "test"))
+                .map(it -> {
+                    Map<String, Object> map = new HashMap<>();
+                    map.put("_key", it.getKey());
+                    map.put("a", it.getAttribute("a"));
+                    map.put("b", it.getAttribute("b"));
+                    return map;
+                })
+                .collect(Collectors.toList());
+
+        final MultiDocumentEntity<DocumentUpdateEntity<BaseDocument>> updateResult = collection
+                .updateDocuments(modifiedDocs, new DocumentUpdateOptions().returnNew(true), BaseDocument.class).get();
+        assertThat(updateResult.getDocuments().size(), is(3));
+        assertThat(updateResult.getErrors().size(), is(0));
+        assertThat(updateResult.getDocuments().stream().map(DocumentUpdateEntity::getNew)
+                        .allMatch(it -> it.getAttribute("a").equals("test") && it.getAttribute("b").equals("test")),
+                is(true));
+    }
+
     @Test
     public void updateDocumentsOne() throws InterruptedException, ExecutionException {
         final Collection<BaseDocument> values = new ArrayList<>();