Fix #97

Note due to: https://jira.mongodb.org/browse/SERVER-9958 This will only work on a mongo>2.5.4, if we want to support<2.5.4 I could try to add detection of the mongo version.
vert-x3 · Dec 21, 2017 · 59faf7b · 59faf7b
1 parent 071b833
commit 59faf7b
Show file tree

Hide file tree

Showing 3 changed files with 125 additions and 4 deletions.
diff --git a/vertx-mongo-client/src/main/java/io/vertx/ext/mongo/impl/MongoClientImpl.java b/vertx-mongo-client/src/main/java/io/vertx/ext/mongo/impl/MongoClientImpl.java
@@ -199,9 +199,13 @@ public io.vertx.ext.mongo.MongoClient updateCollectionWithOptions(String collect
     requireNonNull(options, "options cannot be null");
     requireNonNull(resultHandler, "resultHandler cannot be null");
 
+    update = generateIdIfNeeded(query, update, options);
+
     MongoCollection<JsonObject> coll = getCollection(collection, options.getWriteOption());
     Bson bquery = wrap(encodeKeyWhenUseObjectId(query));
     Bson bupdate = wrap(encodeKeyWhenUseObjectId(update));
+
+
     if (options.isMulti()) {
       coll.updateMany(bquery, bupdate, mongoUpdateOptions(options), toMongoClientUpdateResult(resultHandler));
     } else {
@@ -210,6 +214,23 @@ public io.vertx.ext.mongo.MongoClient updateCollectionWithOptions(String collect
     return this;
   }
 
+  private JsonObject generateIdIfNeeded(JsonObject query, JsonObject update, UpdateOptions options) {
+    if(options.isUpsert() && !update.containsKey(ID_FIELD) && !useObjectId) {
+      JsonObject setId = update.getJsonObject("$setOnInsert", new JsonObject());
+      String id;
+
+      //This seems odd, but if you filter based on _id, mongo expects the generated _id to match
+      if(query.containsKey(ID_FIELD)) {
+        id = query.getString(ID_FIELD);
+      } else {
+        id = JsonObjectCodec.generateHexObjectId();
+      }
+      setId.put(ID_FIELD, id);
+      update.put("$setOnInsert", setId);
+    }
+    return update;
+  }
+
   @Deprecated @Override
   public io.vertx.ext.mongo.MongoClient replace(String collection, JsonObject query, JsonObject replace, Handler<AsyncResult<Void>> resultHandler) {
     replaceWithOptions(collection, query, replace, DEFAULT_UPDATE_OPTIONS, resultHandler);

diff --git a/vertx-mongo-client/src/main/java/io/vertx/ext/mongo/impl/codec/json/JsonObjectCodec.java b/vertx-mongo-client/src/main/java/io/vertx/ext/mongo/impl/codec/json/JsonObjectCodec.java
@@ -42,14 +42,18 @@ public JsonObjectCodec(JsonObject config) {
   public JsonObject generateIdIfAbsentFromDocument(JsonObject json) {
 
     if (!documentHasId(json)) {
-      ObjectId id = new ObjectId();
-
-      if (useObjectId) json.put(ID_FIELD, new JsonObject().put(OID_FIELD, id.toHexString()));
-      else json.put(ID_FIELD, id.toHexString());
+      String value = generateHexObjectId();
+      if (useObjectId) json.put(ID_FIELD, new JsonObject().put(OID_FIELD, value));
+      else json.put(ID_FIELD, value);
     }
     return json;
   }
 
+  public static String generateHexObjectId() {
+    ObjectId id = new ObjectId();
+    return id.toHexString();
+  }
+
   @Override
   public boolean documentHasId(JsonObject json) {
     return json.containsKey(ID_FIELD);

diff --git a/vertx-mongo-client/src/test/java/io/vertx/ext/mongo/MongoClientTest.java b/vertx-mongo-client/src/test/java/io/vertx/ext/mongo/MongoClientTest.java
@@ -1,17 +1,26 @@
 package io.vertx.ext.mongo;
 
+import com.mongodb.async.client.MongoClients;
+import com.mongodb.async.client.MongoDatabase;
 import io.vertx.core.json.JsonObject;
+import io.vertx.ext.mongo.impl.codec.json.JsonObjectCodec;
+import org.junit.After;
+import org.junit.Before;
 import org.junit.Test;
 
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.CountDownLatch;
+import java.util.function.Consumer;
 
 /**
  * @author <a href="http://tfox.org">Tim Fox</a>
  */
 public class MongoClientTest extends MongoClientTestBase {
 
+  private com.mongodb.async.client.MongoClient actualMongo;
+  private MongoDatabase db;
+
   @Override
   public void setUp() throws Exception {
     super.setUp();
@@ -20,11 +29,16 @@ public void setUp() throws Exception {
     CountDownLatch latch = new CountDownLatch(1);
     dropCollections(mongoClient, latch);
     awaitLatch(latch);
+
+
+    actualMongo = MongoClients.create("mongodb://localhost:27018");
+    db = actualMongo.getDatabase(io.vertx.ext.mongo.MongoClient.DEFAULT_DB_NAME);
   }
 
   @Override
   public void tearDown() throws Exception {
     mongoClient.close();
+    actualMongo.close();
     super.tearDown();
   }
 
@@ -50,4 +64,86 @@ public void testFindBatch() throws Exception {
     assertEquals("bar999", foos.get(numDocs - 1));
   }
 
+
+  @Test
+  public void testUpsertCreatesHexIfRecordDoesNotExist() throws Exception {
+    upsertDoc(randomCollection(), createDoc(), null, IGNORE -> {
+      testComplete();
+    });
+
+    await();
+  }
+
+  @Test
+  public void testUpsertWithASetOnInsertIsNotOverWritten() throws Exception {
+    String collection = randomCollection();
+    JsonObject docToInsert = createDoc();
+    JsonObject insertStatement = new JsonObject()
+      .put("$set", docToInsert)
+      .put("$setOnInsert", new JsonObject().put("a-field", "an-entry"));
+
+    upsertDoc(collection, docToInsert, insertStatement, null, saved -> {
+      assertEquals("an-entry", saved.getString("a-field"));
+      testComplete();
+    });
+    await();
+  }
+
+  @Test
+  public void testUpsertDoesNotChangeIdIfRecordExist() throws Exception {
+    String collection = randomCollection();
+    JsonObject docToInsert = createDoc();
+    mongoClient
+      .insert(collection, docToInsert, onSuccess(id -> {
+        upsertDoc(collection, docToInsert, id, IGNORE -> {
+          testComplete();
+        });
+      }));
+    await();
+  }
+
+  private void upsertDoc(String collection, JsonObject docToInsert, String expectedId, Consumer<JsonObject> doneFunction) {
+    JsonObject insertStatement = new JsonObject()
+      .put("$set", docToInsert);
+
+    upsertDoc(collection, docToInsert, insertStatement, expectedId, doneFunction);
+  }
+
+  private void upsertDoc(String collection, JsonObject docToInsert, JsonObject insertStatement, String expectedId, Consumer<JsonObject> doneFunction) {
+    mongoClient.updateCollectionWithOptions(collection,
+      new JsonObject()
+        .put("foo", docToInsert.getString("foo")),
+      insertStatement,
+      new UpdateOptions()
+        .setUpsert(true),
+      onSuccess(res -> {
+        assertEquals(0, res.getDocModified());
+
+        if (expectedId == null) {
+          assertEquals(0, res.getDocMatched());
+          assertNotNull(res.getDocUpsertedId());
+        } else {
+          assertEquals(1, res.getDocMatched());
+          assertNull(res.getDocUpsertedId());
+        }
+
+        //need to check actual DB, not through the Vertx client, in order to make sure the id is a string
+        db
+          .getCollection(collection)
+          .find()
+
+          .first((savedDoc, error) -> {
+            vertx.runOnContext(IGNORE -> {
+              if (expectedId != null) {
+                assertEquals(expectedId, savedDoc.getString(MongoClientUpdateResult.ID_FIELD));
+              } else {
+                assertEquals(res.getDocUpsertedId().getString(MongoClientUpdateResult.ID_FIELD), savedDoc.getString(MongoClientUpdateResult.ID_FIELD));
+              }
+              doneFunction.accept(new JsonObject(savedDoc.toJson()));
+            });
+          });
+      }));
+  }
+
+
 }