feat: adding upsert api with for conditial insert and update

document-store/build.gradle.kts

@@ -8,6 +8,7 @@ plugins {
 
 dependencies {
   api("com.typesafe:config:1.3.2")
+  api("com.google.code.findbugs:jsr305:3.0.2")
   implementation("org.postgresql:postgresql:42.2.13")
   implementation("org.mongodb:mongodb-driver-sync:4.1.2")
   implementation("com.fasterxml.jackson.core:jackson-databind:2.11.0")

document-store/src/main/java/org/hypertrace/core/documentstore/Collection.java

@@ -3,6 +3,7 @@
 import java.io.IOException;
 import java.util.Iterator;
 import java.util.Map;
+import javax.annotation.Nullable;
 
 /** Interface spec for common operations on a collection of documents */
 public interface Collection {
@@ -16,6 +17,21 @@ public interface Collection {
    */
   boolean upsert(Key key, Document document) throws IOException;
 
+  /**
+   * Upsert (create a new doc or update if one already exists) the given document into the doc store
+   * if condition is evaluated to true. Provides optimistic lock support for concurrency update.
+   *
+   * @param key Unique key of the document in the collection.
+   * @param document Document to be upserted.
+   * @param condition Filter condition to be evaluated, on success update the document
+   * @param isUpsert Optional parameter to explicitly control insert or update, default is true.
+   *                 True indicates if the document doesn't exist, it will insert a new document
+   *                 False indicates if the document exists, update it otherwise don't do anything.
+   * @return True if success. False otherwise.
+   */
+  boolean upsert(Key key, Document document, Filter condition, @Nullable Boolean isUpsert)
+      throws IOException;
+
   /**
    * Upsert (create a new doc or update if one already exists) the given document into the doc
    * store.

document-store/src/main/java/org/hypertrace/core/documentstore/mongo/MongoCollection.java

@@ -34,6 +34,7 @@
 import java.util.Set;
 import java.util.function.Function;
 import java.util.stream.Collectors;
+import javax.annotation.Nullable;
 import net.jodah.failsafe.Failsafe;
 import net.jodah.failsafe.RetryPolicy;
 import org.bson.conversions.Bson;
@@ -123,13 +124,36 @@ public boolean upsert(Key key, Document document) throws IOException {
         LOGGER.debug("Write result: " + writeResult.toString());
       }
 
-      return writeResult.getModifiedCount() > 0;
+      return (writeResult.getUpsertedId() != null || writeResult.getModifiedCount() > 0);
     } catch (IOException e) {
       LOGGER.error("Exception upserting document. key: {} content:{}", key, document, e);
       throw e;
     }
   }
 
+  /**
+   * Same as existing upsert method, however, extends the support with condition and optional
+   * parameter for explicitly controlling insert and update.
+   * */
+  @Override
+  public boolean upsert(Key key, Document document, Filter condition, @Nullable Boolean isUpsert)
+      throws IOException {
+    try {
+      Map<String, Object> map = parseQuery(condition);
+      map.put(ID_KEY, key.toString());
+      BasicDBObject conditionObject = new BasicDBObject(map);
+      UpdateOptions options = new UpdateOptions().upsert(isUpsert != null ? isUpsert : true);
+      UpdateResult writeResult =
+          collection.updateOne(conditionObject, this.prepareUpsertWithCondition(key, document), options);
+      if (LOGGER.isDebugEnabled()) {
+        LOGGER.debug("Write result: " + writeResult.toString());
+      }
+      return (writeResult.getUpsertedId() == null || writeResult.getModifiedCount() > 0);
+    } catch (Exception e) {
+      LOGGER.error("Exception upserting document. key: {} content:{}", key, document, e);
+      throw new IOException(e);
+    }
+  }
   /**
    * Adds the following fields automatically: _id, _lastUpdateTime, lastUpdatedTime and created Time
    */
@@ -152,6 +176,21 @@ public Document upsertAndReturn(Key key, Document document) throws IOException {
     return this.dbObjectToDocument(upsertResult);
   }
 
+  private BasicDBObject prepareUpsertWithCondition(Key key, Document document) throws JsonProcessingException {
+    String jsonString = document.toJson();
+    JsonNode jsonNode = MAPPER.readTree(jsonString);
+
+    // escape "." and "$" in field names since Mongo DB does not like them
+    JsonNode sanitizedJsonNode = recursiveClone(jsonNode, this::encodeKey);
+    BasicDBObject setObject = BasicDBObject.parse(MAPPER.writeValueAsString(sanitizedJsonNode));
+    long now = System.currentTimeMillis();
+    setObject.put(ID_KEY, key.toString());
+    setObject.put(LAST_UPDATED_TIME, now);
+    return new BasicDBObject("$set", setObject)
+        .append("$currentDate", new BasicDBObject(LAST_UPDATE_TIME, true))
+        .append("$setOnInsert", new BasicDBObject(CREATED_TIME, now));
+  }
+
   private BasicDBObject prepareUpsert(Key key, Document document) throws JsonProcessingException {
     String jsonString = document.toJson();
     JsonNode jsonNode = MAPPER.readTree(jsonString);