Handled dollar prefixed values for "createOrReplace" Mongo implementa…

…tion (#207)
hypertrace · Sep 9, 2024 · 9aa8fe6 · 9aa8fe6
1 parent d5362f6
commit 9aa8fe6
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 4 deletions.
diff --git a/document-store/src/integrationTest/resources/create/document_one.json b/document-store/src/integrationTest/resources/create/document_one.json
@@ -7,7 +7,7 @@
     },
     {
       "name": "Mars",
-      "color": "Red",
+      "color": "$Red",
       "humans_live": "possibly in future",
       "tags": {}
     }

diff --git a/document-store/src/integrationTest/resources/create/document_one_response.json b/document-store/src/integrationTest/resources/create/document_one_response.json
@@ -8,7 +8,7 @@
       },
       {
         "name": "Mars",
-        "color": "Red",
+        "color": "$Red",
         "humans_live": "possibly in future",
         "tags": {}
       }

diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/MongoUtils.java b/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/MongoUtils.java
@@ -14,8 +14,10 @@
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.ArrayNode;
 import com.fasterxml.jackson.databind.node.JsonNodeFactory;
 import com.fasterxml.jackson.databind.node.ObjectNode;
+import com.fasterxml.jackson.databind.node.TextNode;
 import com.mongodb.BasicDBObject;
 import com.mongodb.client.model.ReturnDocument;
 import java.io.IOException;
@@ -35,6 +37,7 @@
 public final class MongoUtils {
   public static final String FIELD_SEPARATOR = ".";
   public static final String PREFIX = "$";
+  private static final String UNICODE_FOR_FIELD_PREFIX = "\\u0024";
   private static final String LITERAL = PREFIX + "literal";
   private static final String UNSUPPORTED_OPERATION = "No MongoDB support available for: '%s'";
   private static final ObjectMapper MAPPER = new ObjectMapper();
@@ -54,15 +57,19 @@ public static String encodeKey(final String key) {
       return null;
     }
 
-    return key.replace("\\", "\\\\").replace(PREFIX, "\\u0024").replace(FIELD_SEPARATOR, "\\u002e");
+    return key.replace("\\", "\\\\")
+        .replace(PREFIX, UNICODE_FOR_FIELD_PREFIX)
+        .replace(FIELD_SEPARATOR, "\\u002e");
   }
 
   public static String decodeKey(final String key) {
     if (key == null) {
       return null;
     }
 
-    return key.replace("\\u002e", FIELD_SEPARATOR).replace("\\u0024", PREFIX).replace("\\\\", "\\");
+    return key.replace("\\u002e", FIELD_SEPARATOR)
+        .replace(UNICODE_FOR_FIELD_PREFIX, PREFIX)
+        .replace("\\\\", "\\");
   }
 
   public static String getLastField(final String fieldPath) {
@@ -99,6 +106,10 @@ public static JsonNode recursiveClone(
       JsonNode src,
       UnaryOperator<String> function,
       final UnaryOperator<ObjectNode> emptyObjectConverter) {
+    if (src.isTextual()) {
+      return new TextNode(function.apply(src.asText()));
+    }
+
     if (!src.isObject()) {
       return src;
     }
@@ -113,6 +124,17 @@ public static JsonNode recursiveClone(
       if (value.isObject()) {
         newValue = recursiveClone(value, function, emptyObjectConverter);
       }
+      if (value.isArray()) {
+        newValue = new ArrayNode(JsonNodeFactory.instance);
+        for (int i = 0; i < value.size(); i++) {
+          final JsonNode elementValue =
+              recursiveClone(value.get(i), function, emptyObjectConverter);
+          ((ArrayNode) newValue).add(elementValue);
+        }
+      }
+      if (newValue.isTextual()) {
+        newValue = new TextNode(function.apply(newValue.asText()));
+      }
       tgt.set(newFieldName, newValue);
     }
     return tgt.isEmpty() ? emptyObjectConverter.apply(tgt) : tgt;