perf: optimize mongo update impl when return document type is none (#241)

* perf: optimize mongo update impl when return document type is none

* test: fix unit test

* test: add integration test

* Revert "test: fix unit test"

This reverts commit a01a694.

* Revert "perf: optimize mongo update impl when return document type is none"

This reverts commit d85dc00.

* perf: optimize update impl for mongo by projecting only _id field when return document types is none

Author: GurtejSohi

document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreQueryV1Test.java

@@ -2510,6 +2510,107 @@ public void testAtomicUpdateDocumentWithoutSelections(final String datastoreName
           "query/updatable_collection_data_without_selection.json",
           9);
     }
+
+    @ParameterizedTest
+    @ArgumentsSource(AllProvider.class)
+    public void testAtomicUpdateWithReturnDocumentTypeNone(final String datastoreName)
+        throws IOException {
+      final Collection collection = getCollection(datastoreName, UPDATABLE_COLLECTION_NAME);
+      createCollectionData("query/updatable_collection_data.json", UPDATABLE_COLLECTION_NAME);
+
+      final Query query =
+          Query.builder()
+              .setFilter(
+                  LogicalExpression.builder()
+                      .operator(AND)
+                      .operand(
+                          RelationalExpression.of(
+                              IdentifierExpression.of("item"), EQ, ConstantExpression.of("Soap")))
+                      .operand(
+                          RelationalExpression.of(
+                              IdentifierExpression.of("date"),
+                              LT,
+                              ConstantExpression.of("2022-08-09T18:53:17Z")))
+                      .build())
+              .addSort(SortingSpec.of(IdentifierExpression.of("price"), ASC))
+              .addSort(SortingSpec.of(IdentifierExpression.of("date"), DESC))
+              .addSelection(IdentifierExpression.of("quantity"))
+              .addSelection(IdentifierExpression.of("price"))
+              .addSelection(IdentifierExpression.of("date"))
+              .addSelection(IdentifierExpression.of("props"))
+              .build();
+
+      final SubDocumentUpdate dateUpdate = SubDocumentUpdate.of("date", "2022-08-09T18:53:17Z");
+      final SubDocumentUpdate quantityUpdate = SubDocumentUpdate.of("quantity", 1000);
+      final SubDocumentUpdate propsUpdate =
+          SubDocumentUpdate.of(
+              "props", SubDocumentValue.of(new JSONDocument("{\"brand\": \"Dettol\"}")));
+
+      // Perform update with ReturnDocumentType.NONE
+      final Optional<Document> updateResult =
+          collection.update(
+              query,
+              List.of(dateUpdate, quantityUpdate, propsUpdate),
+              UpdateOptions.builder().returnDocumentType(NONE).build());
+
+      // Verify that no document is returned (as expected with NONE)
+      assertFalse(
+          updateResult.isPresent(), "Should return empty Optional when ReturnDocumentType.NONE");
+
+      // Verify that the document was actually updated in the database by querying it
+      final Query verificationQuery =
+          Query.builder()
+              .setFilter(
+                  LogicalExpression.builder()
+                      .operator(AND)
+                      .operand(
+                          RelationalExpression.of(
+                              IdentifierExpression.of("item"), EQ, ConstantExpression.of("Soap")))
+                      .operand(
+                          RelationalExpression.of(
+                              IdentifierExpression.of("date"),
+                              EQ,
+                              ConstantExpression.of("2022-08-09T18:53:17Z")))
+                      .build())
+              .build();
+
+      final CloseableIterator<Document> verificationResults = collection.find(verificationQuery);
+      assertTrue(verificationResults.hasNext(), "Updated document should exist in database");
+
+      final Document updatedDocument = verificationResults.next();
+      final JsonNode json = new ObjectMapper().readTree(updatedDocument.toJson());
+      // Verify that the right document was updated in the database
+      assertEquals(null, json.get("sales"), "'sales' field should not be present");
+
+      // Verify the fields were updated
+      assertEquals(1000, json.get("quantity").asInt(), "Quantity should be updated to 1000");
+      assertEquals("2022-08-09T18:53:17Z", json.get("date").asText(), "Date should be updated");
+      final JsonNode props = json.get("props");
+      assertEquals(
+          "Dettol", props.get("brand").asText(), "Props brand should be updated to Dettol");
+
+      verificationResults.close();
+
+      // Additional test: Verify that the same update with AFTER_UPDATE returns the document
+      // Reset the data first
+      createCollectionData("query/updatable_collection_data.json", UPDATABLE_COLLECTION_NAME);
+
+      final Optional<Document> afterUpdateResult =
+          collection.update(
+              query,
+              List.of(dateUpdate, quantityUpdate, propsUpdate),
+              UpdateOptions.builder().returnDocumentType(AFTER_UPDATE).build());
+
+      // This should return the updated document
+      assertTrue(
+          afterUpdateResult.isPresent(),
+          "Should return document when ReturnDocumentType.AFTER_UPDATE");
+
+      final Document returnedDocument = afterUpdateResult.get();
+      final JsonNode returnedJson = new ObjectMapper().readTree(returnedDocument.toJson());
+      assertEquals(1000, returnedJson.get("quantity").asInt());
+      assertEquals("2022-08-09T18:53:17Z", returnedJson.get("date").asText());
+    }
   }
 
   @Nested

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

@@ -2,6 +2,7 @@
 
 import static org.hypertrace.core.documentstore.model.options.DataFreshness.REALTIME_FRESHNESS;
 import static org.hypertrace.core.documentstore.model.options.ReturnDocumentType.NONE;
+import static org.hypertrace.core.documentstore.mongo.MongoCollection.ID_KEY;
 import static org.hypertrace.core.documentstore.mongo.MongoUtils.getReturnDocument;
 import static org.hypertrace.core.documentstore.mongo.query.parser.MongoFilterTypeExpressionParser.getFilter;
 import static org.hypertrace.core.documentstore.mongo.query.parser.MongoSelectTypeExpressionParser.getSelections;
@@ -54,7 +55,6 @@ public Optional<Document> update(
     updateValidator.validate(updates);
 
     try {
-      final BasicDBObject selections = getSelections(query);
       final BasicDBObject sorts = getOrders(query);
       final FindOneAndUpdateOptions options = new FindOneAndUpdateOptions();
       options.upsert(
@@ -65,8 +65,14 @@ public Optional<Document> update(
 
       options.returnDocument(getReturnDocument(returnDocumentType));
 
-      if (!selections.isEmpty()) {
-        options.projection(selections);
+      if (returnDocumentType == NONE) {
+        // Optimize by projecting only the _id field when we don't need to return the document
+        options.projection(new BasicDBObject(ID_KEY, 1));
+      } else {
+        final BasicDBObject selections = getSelections(query);
+        if (!selections.isEmpty()) {
+          options.projection(selections);
+        }
       }
 
       if (!sorts.isEmpty()) {

document-store/src/test/java/org/hypertrace/core/documentstore/mongo/MongoCollectionTest.java

@@ -322,7 +322,7 @@ void testUpdateAtomicWithFilter_getNone() throws IOException {
 
       assertFalse(result.isPresent());
 
-      assertEquals(selections, options.getValue().getProjection());
+      assertEquals(new BasicDBObject("_id", 1), options.getValue().getProjection());
       assertEquals(sort, options.getValue().getSort());
       assertEquals(AFTER, options.getValue().getReturnDocument());