exposed update options to client

hypertrace · Aug 16, 2024 · 194f5e5 · 194f5e5
1 parent 4f62e48
commit 194f5e5
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 7 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,14 +2,19 @@
 
 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).build();
+      UpdateOptions.builder().returnDocumentType(AFTER_UPDATE).updateOptions(new com.mongodb.client.model.UpdateOptions())
+          .findOneAndUpdateOptions(new FindOneAndUpdateOptions())
+          .build();
 
   ReturnDocumentType returnDocumentType;
+  com.mongodb.client.model.UpdateOptions updateOptions;
+  FindOneAndUpdateOptions findOneAndUpdateOptions;
 }
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,7 @@ public Optional<Document> update(
     try {
       final BasicDBObject selections = getSelections(query);
       final BasicDBObject sorts = getOrders(query);
-      final FindOneAndUpdateOptions options = new FindOneAndUpdateOptions();
+      final FindOneAndUpdateOptions options = updateOptions.getFindOneAndUpdateOptions();
       final ReturnDocumentType returnDocumentType = updateOptions.getReturnDocumentType();
 
       options.returnDocument(getReturnDocument(returnDocumentType));
@@ -96,29 +96,29 @@ public Optional<MongoCursor<BasicDBObject>> bulkUpdate(
     switch (returnDocumentType) {
       case BEFORE_UPDATE:
         cursor = queryExecutor.aggregate(query);
-        logAndUpdate(filter, updateObject);
+        logAndUpdate(filter, updateObject, updateOptions.getUpdateOptions());
         return Optional.of(cursor);
 
       case AFTER_UPDATE:
-        logAndUpdate(filter, updateObject);
+        logAndUpdate(filter, updateObject, updateOptions.getUpdateOptions());
         cursor = queryExecutor.aggregate(query);
         return Optional.of(cursor);
 
       case NONE:
-        logAndUpdate(filter, updateObject);
+        logAndUpdate(filter, updateObject, updateOptions.getUpdateOptions());
         return Optional.empty();
 
       default:
         throw new IOException("Unrecognized return document type: " + returnDocumentType);
     }
   }
 
-  private void logAndUpdate(final BasicDBObject filter, final BasicDBObject setObject)
+  private void logAndUpdate(final BasicDBObject filter, final BasicDBObject setObject, com.mongodb.client.model.UpdateOptions updateOptions)
       throws IOException {
     try {
       log.debug(
           "Updating {} using {} with filter {}", collection.getNamespace(), setObject, filter);
-      collection.updateMany(filter, setObject);
+      collection.updateMany(filter, setObject, updateOptions);
     } catch (Exception e) {
       throw new IOException("Error while updating", e);
     }