From e7ad62d84a9b8e210b673c456c22bf02ea230a28 Mon Sep 17 00:00:00 2001
From: shubhamgupta <shubham.gupta@traceable.ai>
Date: Wed, 21 Aug 2024 13:58:08 +0530
Subject: [PATCH] after review

---
 .../documentstore/model/options/UpdateOptions.java   | 10 ++--------
 .../mongo/update/MongoUpdateExecutor.java            | 12 ++++++++----
 2 files changed, 10 insertions(+), 12 deletions(-)

diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/model/options/UpdateOptions.java b/document-store/src/main/java/org/hypertrace/core/documentstore/model/options/UpdateOptions.java
index ad9a29f3..1017171e 100644
--- a/document-store/src/main/java/org/hypertrace/core/documentstore/model/options/UpdateOptions.java
+++ b/document-store/src/main/java/org/hypertrace/core/documentstore/model/options/UpdateOptions.java
@@ -2,7 +2,6 @@
 
 import static org.hypertrace.core.documentstore.model.options.ReturnDocumentType.AFTER_UPDATE;
 
-import com.mongodb.client.model.FindOneAndUpdateOptions;
 import lombok.Builder;
 import lombok.Value;
 
@@ -10,13 +9,8 @@
 @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/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/update/MongoUpdateExecutor.java b/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/update/MongoUpdateExecutor.java
index fd68723c..15109603 100644
--- a/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/update/MongoUpdateExecutor.java
+++ b/document-store/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: