fix(patch): consider escaped characters when applying JSON patches (d…

…atahub-project#10717) Co-authored-by: John Joyce <john@acryl.io>
dushayntAW · Jul 1, 2024 · 8edc94d · 8edc94d
1 parent 93616f7
commit 8edc94d
Show file tree

Hide file tree

Showing 2 changed files with 67 additions and 1 deletion.
diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/patch/template/TemplateUtil.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/patch/template/TemplateUtil.java
@@ -85,12 +85,18 @@ public static JsonNode populateTopLevelKeys(JsonNode transformedNode, JsonPatch
       // Skip first as it will always be blank due to path starting with /
       for (int i = 1; i < endIdx; i++) {
         if (parent.get(keys[i]) == null) {
-          ((ObjectNode) parent).set(keys[i], instance.objectNode());
+          String decodedKey = decodeValue(keys[i]);
+          ((ObjectNode) parent).set(decodedKey, instance.objectNode());
         }
         parent = parent.get(keys[i]);
       }
     }
 
     return transformedNodeClone;
   }
+
+  /** Simply decode a JSON-patch encoded value * */
+  private static String decodeValue(String value) {
+    return value.replace("~1", "/").replace("~0", "~");
+  }
 }
diff --git a/...rc/test/java/com/linkedin/metadata/aspect/patch/template/UpstreamLineageTemplateTest.java b/...rc/test/java/com/linkedin/metadata/aspect/patch/template/UpstreamLineageTemplateTest.java
@@ -4,6 +4,7 @@
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertTrue;
 
+import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.linkedin.common.UrnArray;
 import com.linkedin.common.urn.Urn;
@@ -254,4 +255,63 @@ public void testLargePatchStandard() throws Exception {
     assertEquals(result.getUpstreams().size(), 187, "Expected 1 less upstream");
     assertEquals(result.getFineGrainedLineages().size(), 607);
   }
+
+  @Test
+  public void testPatchWithFieldWithForwardSlash() throws JsonProcessingException {
+
+    String downstreamUrn =
+        "/fineGrainedLineages/CREATE/urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:bigquery,upstream_table_1,PROD),c1)";
+    String unescapedUpstreamUrn =
+        "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:bigquery,upstream_table_2,PROD),slash/column)";
+    String escapedUpstreamUrn =
+        "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:bigquery,upstream_table_2,PROD),slash~1column)";
+    String lineagePath = downstreamUrn + "//" + escapedUpstreamUrn;
+
+    UpstreamLineageTemplate upstreamLineageTemplate = new UpstreamLineageTemplate();
+    UpstreamLineage upstreamLineage = upstreamLineageTemplate.getDefault();
+    JsonPatchBuilder jsonPatchBuilder = Json.createPatchBuilder();
+
+    JsonObjectBuilder fineGrainedLineageNode = Json.createObjectBuilder();
+    JsonValue upstreamConfidenceScore = Json.createValue(1.0f);
+    fineGrainedLineageNode.add("confidenceScore", upstreamConfidenceScore);
+
+    jsonPatchBuilder.add(lineagePath, fineGrainedLineageNode.build());
+
+    // Initial population test
+    UpstreamLineage result =
+        upstreamLineageTemplate.applyPatch(upstreamLineage, jsonPatchBuilder.build());
+
+    assertEquals(
+        result.getFineGrainedLineages().get(0).getUpstreams().get(0).toString(),
+        unescapedUpstreamUrn);
+  }
+
+  @Test
+  public void testPatchWithFieldWithTilde() throws JsonProcessingException {
+
+    String downstreamUrn =
+        "/fineGrainedLineages/CREATE/urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:bigquery,upstream_table_1,PROD),c1)";
+    String unescapedUpstreamUrn =
+        "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:bigquery,upstream_table_2,PROD),tilde~column)";
+    String escapedUpstreamUrn =
+        "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:bigquery,upstream_table_2,PROD),tilde~0column)";
+    String lineagePath = downstreamUrn + "//" + escapedUpstreamUrn;
+
+    UpstreamLineageTemplate upstreamLineageTemplate = new UpstreamLineageTemplate();
+    UpstreamLineage upstreamLineage = upstreamLineageTemplate.getDefault();
+    JsonPatchBuilder jsonPatchBuilder = Json.createPatchBuilder();
+
+    JsonObjectBuilder fineGrainedLineageNode = Json.createObjectBuilder();
+    JsonValue upstreamConfidenceScore = Json.createValue(1.0f);
+    fineGrainedLineageNode.add("confidenceScore", upstreamConfidenceScore);
+
+    jsonPatchBuilder.add(lineagePath, fineGrainedLineageNode.build());
+
+    // Initial population test
+    UpstreamLineage result =
+        upstreamLineageTemplate.applyPatch(upstreamLineage, jsonPatchBuilder.build());
+    assertEquals(
+        result.getFineGrainedLineages().get(0).getUpstreams().get(0).toString(),
+        unescapedUpstreamUrn);
+  }
 }