Fix #219: add $setOnInsert update operation

stargate · Mar 6, 2023 · 9099b23 · 9099b23
1 parent a6ea098
commit 9099b23
Show file tree

Hide file tree

Showing 7 changed files with 218 additions and 9 deletions.
diff --git a/...n/java/io/stargate/sgv2/jsonapi/api/model/command/clause/update/SetOnInsertOperation.java b/...n/java/io/stargate/sgv2/jsonapi/api/model/command/clause/update/SetOnInsertOperation.java
@@ -0,0 +1,80 @@
+package io.stargate.sgv2.jsonapi.api.model.command.clause.update;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import io.stargate.sgv2.jsonapi.util.JsonUtil;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Objects;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+/** Implementation of {@code $set} update operation used to assign values to document fields. */
+public class SetOnInsertOperation extends UpdateOperation {
+  private final List<SetOnInsertAction> actions;
+
+  private SetOnInsertOperation(List<SetOnInsertAction> actions) {
+    this.actions = sortByPath(actions);
+  }
+
+  public static SetOnInsertOperation construct(ObjectNode args) {
+    List<SetOnInsertAction> additions = new ArrayList<>();
+    var it = args.fields();
+    while (it.hasNext()) {
+      var entry = it.next();
+      String path = validateUpdatePath(UpdateOperator.SET, entry.getKey());
+      additions.add(new SetOnInsertAction(path, entry.getValue()));
+    }
+    return new SetOnInsertOperation(additions);
+  }
+
+  /**
+   * Override method used to set update filter condition fields to the document
+   *
+   * @param filterPath
+   * @param value
+   * @return
+   */
+  public static SetOnInsertOperation construct(String filterPath, JsonNode value) {
+    List<SetOnInsertAction> additions = new ArrayList<>();
+    String path = validateUpdatePath(UpdateOperator.SET_ON_INSERT, filterPath);
+    additions.add(new SetOnInsertAction(path, value));
+    return new SetOnInsertOperation(additions);
+  }
+
+  @Override
+  public boolean shouldApplyIf(boolean isInsert) {
+    // Only run for true inserts; skip otherwise
+    return isInsert;
+  }
+
+  @Override
+  public boolean updateDocument(ObjectNode doc, UpdateTargetLocator targetLocator) {
+    boolean modified = false;
+    for (SetOnInsertAction addition : actions) {
+      UpdateTarget target = targetLocator.findOrCreate(doc, addition.path());
+      JsonNode newValue = addition.value();
+      JsonNode oldValue = target.valueNode();
+
+      // Modify if no old value OR new value differs, as per Mongo-equality rules
+      if ((oldValue == null) || !JsonUtil.equalsOrdered(oldValue, newValue)) {
+        target.replaceValue(newValue);
+        modified = true;
+      }
+    }
+    return modified;
+  }
+
+  public Set<String> getPaths() {
+    return actions.stream().map(SetOnInsertAction::path).collect(Collectors.toSet());
+  }
+
+  // Just needed for tests
+  @Override
+  public boolean equals(Object o) {
+    return (o instanceof SetOnInsertOperation)
+        && Objects.equals(this.actions, ((SetOnInsertOperation) o).actions);
+  }
+
+  private record SetOnInsertAction(String path, JsonNode value) implements ActionWithPath {}
+}
diff --git a/src/main/java/io/stargate/sgv2/jsonapi/api/model/command/clause/update/UpdateOperation.java b/src/main/java/io/stargate/sgv2/jsonapi/api/model/command/clause/update/UpdateOperation.java
@@ -21,6 +21,19 @@ public abstract class UpdateOperation {
    */
   public abstract boolean updateDocument(ObjectNode doc, UpdateTargetLocator targetLocator);
 
+  /**
+   * Method called to see if update operator should be applied for specific kind of update:
+   * currently one special case is that of document insertion as part of upsert. Most update
+   * operations should apply for all updates so the default implementation returns {@code true};
+   *
+   * @param isInsert True if the document to update was just inserted (as part of upsert operation)
+   * @return {@code true} If the update should be applied for document context; {@code false} if it
+   *     should be skipped
+   */
+  public boolean shouldApplyIf(boolean isInsert) {
+    return true;
+  }
+
   /**
    * Shared validation method used by mutating operations (like {@code $set}, {@code $unset}, {@code
    * inc}, {@code pop}) to ensure they are not used to modify paths that are not allowed:

diff --git a/src/main/java/io/stargate/sgv2/jsonapi/api/model/command/clause/update/UpdateOperator.java b/src/main/java/io/stargate/sgv2/jsonapi/api/model/command/clause/update/UpdateOperator.java
@@ -42,6 +42,13 @@ public UpdateOperation resolveOperation(ObjectNode arguments) {
       return SetOperation.construct(arguments);
     }
   },
+  SET_ON_INSERT("$setOnInsert") {
+    @Override
+    public UpdateOperation resolveOperation(ObjectNode arguments) {
+      return SetOnInsertOperation.construct(arguments);
+    }
+  },
+
   UNSET("$unset") {
     @Override
     public UpdateOperation resolveOperation(ObjectNode arguments) {

diff --git a/...in/java/io/stargate/sgv2/jsonapi/service/operation/model/impl/ReadAndUpdateOperation.java b/...in/java/io/stargate/sgv2/jsonapi/service/operation/model/impl/ReadAndUpdateOperation.java
@@ -90,10 +90,11 @@ public Uni<Supplier<CommandResult>> execute(QueryExecutor queryExecutor) {
             .onItem()
             .transformToUniAndConcatenate(
                 readDocument -> {
-                  DocumentUpdater.DocumentUpdaterResponse documentUpdaterResponse =
-                      documentUpdater().applyUpdates(readDocument.document().deepCopy());
                   final JsonNode originalDocument =
                       readDocument.txnId() == null ? null : readDocument.document();
+                  final boolean isInsert = (originalDocument == null);
+                  DocumentUpdater.DocumentUpdaterResponse documentUpdaterResponse =
+                      documentUpdater().applyUpdates(readDocument.document().deepCopy(), isInsert);
                   JsonNode updatedDocument = documentUpdaterResponse.document();
                   Uni<DocumentId> updated = Uni.createFrom().nullItem();
                   if (documentUpdaterResponse.modified()) {

diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/updater/DocumentUpdater.java b/src/main/java/io/stargate/sgv2/jsonapi/service/updater/DocumentUpdater.java
@@ -13,12 +13,20 @@ public static DocumentUpdater construct(UpdateClause updateDef) {
     return new DocumentUpdater(updateDef.buildOperations());
   }
 
-  public DocumentUpdaterResponse applyUpdates(JsonNode readDocument) {
+  /**
+   * @param readDocument Document to update
+   * @param docInserted True if document was just created (inserted); false if updating existing
+   *     document
+   */
+  public DocumentUpdaterResponse applyUpdates(JsonNode readDocument, boolean docInserted) {
     UpdateTargetLocator targetLocator = new UpdateTargetLocator();
     ObjectNode docToUpdate = (ObjectNode) readDocument;
     boolean modified = false;
-    for (UpdateOperation updateOperation : updateOperations)
-      modified |= updateOperation.updateDocument(docToUpdate, targetLocator);
+    for (UpdateOperation updateOperation : updateOperations) {
+      if (updateOperation.shouldApplyIf(docInserted)) {
+        modified |= updateOperation.updateDocument(docToUpdate, targetLocator);
+      }
+    }
     return new DocumentUpdaterResponse(readDocument, modified);
   }
 

diff --git a/src/test/java/io/stargate/sgv2/jsonapi/api/v1/FindAndUpdateIntegrationTest.java b/src/test/java/io/stargate/sgv2/jsonapi/api/v1/FindAndUpdateIntegrationTest.java
@@ -6,6 +6,7 @@
 import static org.hamcrest.Matchers.emptyString;
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.not;
+import static org.hamcrest.Matchers.nullValue;
 
 import io.quarkus.test.common.QuarkusTestResource;
 import io.quarkus.test.junit.QuarkusIntegrationTest;
@@ -1683,6 +1684,105 @@ public void findByColumnAndAddToSetWithEach() {
     }
   }
 
+  @Nested
+  class UpdateOneWithSetOnInsert {
+    @Test
+    @Order(2)
+    public void findByIdUpsertAndAddOnInsert() {
+      String json =
+          """
+          {
+            "findOneAndUpdate": {
+              "filter" : {"_id" : "setOnInsertDoc1"},
+              "update" : {
+                  "$set" : {"active_user": true},
+                  "$setOnInsert" : {"new_user": true}
+              },
+              "options" : {"returnDocument" : "after", "upsert" : true}
+            }
+          }
+          """;
+      // On Insert (for Upsert) should apply both $set and $setOnInsert
+      String expected = "{\"_id\":\"setOnInsertDoc1\", \"new_user\":true, \"active_user\":true}";
+      given()
+          .header(HttpConstants.AUTHENTICATION_TOKEN_HEADER_NAME, getAuthToken())
+          .contentType(ContentType.JSON)
+          .body(json)
+          .when()
+          .post(CollectionResource.BASE_PATH, keyspaceId.asInternal(), collectionName)
+          .then()
+          .statusCode(200)
+          .body("data.docs[0]", jsonEquals(expected))
+          .body("status.upsertedId", is("setOnInsertDoc1"))
+          .body("status.matchedCount", is(0))
+          .body("status.modifiedCount", is(0));
+
+      given()
+          .header(HttpConstants.AUTHENTICATION_TOKEN_HEADER_NAME, getAuthToken())
+          .contentType(ContentType.JSON)
+          .body(
+              """
+                  {
+                    "find": {
+                      "filter" : {"_id" : "setOnInsertDoc1"}
+                    }
+                  }
+                  """)
+          .when()
+          .post(CollectionResource.BASE_PATH, keyspaceId.asInternal(), collectionName)
+          .then()
+          .statusCode(200)
+          .body("data.docs[0]", jsonEquals(expected));
+
+      // However: with update for upsert, $setOnInsert not to be applied
+      json =
+          """
+            {
+              "findOneAndUpdate": {
+                "filter" : {"_id" : "setOnInsertDoc1"},
+                "update" : {
+                    "$set" : {"new_user": false},
+                    "$setOnInsert" : {"x": 5}
+                },
+                "options" : {"returnDocument" : "after", "upsert" : true}
+              }
+            }
+            """;
+      expected = "{\"_id\":\"setOnInsertDoc1\", \"new_user\":false, \"active_user\":true}";
+      given()
+          .header(HttpConstants.AUTHENTICATION_TOKEN_HEADER_NAME, getAuthToken())
+          .contentType(ContentType.JSON)
+          .body(json)
+          .when()
+          .post(CollectionResource.BASE_PATH, keyspaceId.asInternal(), collectionName)
+          .then()
+          .statusCode(200)
+          .body("data.docs[0]", jsonEquals(expected))
+          .body("status.matchedCount", is(1))
+          .body("status.modifiedCount", is(1))
+          .body("status.upsertedId", nullValue());
+
+      // And validate to make sure nothing was actually modified
+      json =
+          """
+            {
+              "find": {
+                "filter" : {"_id" : "setOnInsertDoc1"}
+              }
+            }
+            """;
+      given()
+          .header(HttpConstants.AUTHENTICATION_TOKEN_HEADER_NAME, getAuthToken())
+          .contentType(ContentType.JSON)
+          .body(json)
+          .when()
+          .post(CollectionResource.BASE_PATH, keyspaceId.asInternal(), collectionName)
+          .then()
+          .statusCode(200)
+          .body("data.docs[0]", jsonEquals(expected));
+    }
+  }
+
   private void insertDoc(String docJson) {
     String doc =
         """

diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/updater/DocumentUpdaterTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/updater/DocumentUpdaterTest.java
@@ -52,7 +52,7 @@ public void setUpdateCondition() throws Exception {
                   UpdateOperator.SET,
                   objectMapper.getNodeFactory().objectNode().put("location", "New York")));
       DocumentUpdater.DocumentUpdaterResponse updatedDocument =
-          documentUpdater.applyUpdates(baseData);
+          documentUpdater.applyUpdates(baseData, false);
       assertThat(updatedDocument)
           .isNotNull()
           .satisfies(
@@ -81,7 +81,7 @@ public void setUpdateNewData() throws Exception {
                   UpdateOperator.SET,
                   objectMapper.getNodeFactory().objectNode().put("new_data", "data")));
       DocumentUpdater.DocumentUpdaterResponse updatedDocument =
-          documentUpdater.applyUpdates(baseData);
+          documentUpdater.applyUpdates(baseData, false);
       assertThat(updatedDocument)
           .isNotNull()
           .satisfies(
@@ -110,7 +110,7 @@ public void setUpdateNumberData() throws Exception {
                   UpdateOperator.SET,
                   objectMapper.getNodeFactory().objectNode().put("new_data", 40)));
       DocumentUpdater.DocumentUpdaterResponse updatedDocument =
-          documentUpdater.applyUpdates(baseData);
+          documentUpdater.applyUpdates(baseData, false);
       assertThat(updatedDocument)
           .isNotNull()
           .satisfies(
@@ -138,7 +138,7 @@ public void unsetUpdateData() throws Exception {
               DocumentUpdaterUtils.updateClause(
                   UpdateOperator.UNSET, objectMapper.getNodeFactory().objectNode().put("col", 1)));
       DocumentUpdater.DocumentUpdaterResponse updatedDocument =
-          documentUpdater.applyUpdates(baseData);
+          documentUpdater.applyUpdates(baseData, false);
       assertThat(updatedDocument)
           .isNotNull()
           .satisfies(