after review

hypertrace · Aug 21, 2024 · e7ad62d · e7ad62d
1 parent 934831b
commit e7ad62d
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 12 deletions.
diff --git a/...nt-store/src/main/java/org/hypertrace/core/documentstore/model/options/UpdateOptions.java b/...nt-store/src/main/java/org/hypertrace/core/documentstore/model/options/UpdateOptions.java
@@ -2,21 +2,15 @@
 
 import static org.hypertrace.core.documentstore.model.options.ReturnDocumentType.AFTER_UPDATE;
 
-import com.mongodb.client.model.FindOneAndUpdateOptions;
 import lombok.Builder;
 import lombok.Value;
 
 @Value
 @Builder
 public class UpdateOptions {
   public static UpdateOptions DEFAULT_UPDATE_OPTIONS =
-      UpdateOptions.builder()
-          .returnDocumentType(AFTER_UPDATE)
-          .updateOptions(new com.mongodb.client.model.UpdateOptions())
-          .findOneAndUpdateOptions(new FindOneAndUpdateOptions())
-          .build();
+      UpdateOptions.builder().returnDocumentType(AFTER_UPDATE).upsert(false).build();
 
   ReturnDocumentType returnDocumentType;
-  com.mongodb.client.model.UpdateOptions updateOptions;
-  FindOneAndUpdateOptions findOneAndUpdateOptions;
+  boolean upsert;
 }
diff --git a/...ore/src/main/java/org/hypertrace/core/documentstore/mongo/update/MongoUpdateExecutor.java b/...ore/src/main/java/org/hypertrace/core/documentstore/mongo/update/MongoUpdateExecutor.java
@@ -53,7 +53,8 @@ public Optional<Document> update(
     try {
       final BasicDBObject selections = getSelections(query);
       final BasicDBObject sorts = getOrders(query);
-      final FindOneAndUpdateOptions options = updateOptions.getFindOneAndUpdateOptions();
+      final FindOneAndUpdateOptions options = new FindOneAndUpdateOptions();
+      options.upsert(updateOptions.isUpsert());
       final ReturnDocumentType returnDocumentType = updateOptions.getReturnDocumentType();
 
       options.returnDocument(getReturnDocument(returnDocumentType));
@@ -91,21 +92,24 @@ public Optional<MongoCursor<BasicDBObject>> bulkUpdate(
     final BasicDBObject filter = getFilter(query, Query::getFilter);
     final BasicDBObject updateObject = updateParser.buildUpdateClause(updates);
     final ReturnDocumentType returnDocumentType = updateOptions.getReturnDocumentType();
+    final com.mongodb.client.model.UpdateOptions mongoUpdateOptions =
+        new com.mongodb.client.model.UpdateOptions();
+    mongoUpdateOptions.upsert(updateOptions.isUpsert());
     final MongoCursor<BasicDBObject> cursor;
 
     switch (returnDocumentType) {
       case BEFORE_UPDATE:
         cursor = queryExecutor.aggregate(query);
-        logAndUpdate(filter, updateObject, updateOptions.getUpdateOptions());
+        logAndUpdate(filter, updateObject, mongoUpdateOptions);
         return Optional.of(cursor);
 
       case AFTER_UPDATE:
-        logAndUpdate(filter, updateObject, updateOptions.getUpdateOptions());
+        logAndUpdate(filter, updateObject, mongoUpdateOptions);
         cursor = queryExecutor.aggregate(query);
         return Optional.of(cursor);
 
       case NONE:
-        logAndUpdate(filter, updateObject, updateOptions.getUpdateOptions());
+        logAndUpdate(filter, updateObject, mongoUpdateOptions);
         return Optional.empty();
 
       default: