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

datahub-project · Jun 16, 2024 · 33142d1 · 33142d1
1 parent 402bf31
commit 33142d1
Show file tree

Hide file tree

Showing 2 changed files with 64 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
@@ -13,6 +13,7 @@
 import jakarta.json.JsonPatch;
 import jakarta.json.JsonValue;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 
 public class TemplateUtil {
@@ -73,7 +74,11 @@ public static JsonNode populateTopLevelKeys(JsonNode transformedNode, JsonPatch
     JsonNode transformedNodeClone = transformedNode.deepCopy();
     List<Pair<PatchOperationType, String>> paths = getPaths(jsonPatch);
     for (Pair<PatchOperationType, String> operationPath : paths) {
-      String[] keys = operationPath.getSecond().split("/");
+      // a JsonPatch will have `~` encoded as `~0` and `/` encoded as `~1`
+      String[] keys =
+          Arrays.stream(operationPath.getSecond().split("/"))
+              .map(k -> k.replace("~0", "~").replace("~1", "/"))
+              .toArray(String[]::new);
       JsonNode parent = transformedNodeClone;
 
       // if not remove, skip last key as we only need to populate top level

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,61 @@ 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);
+  }
 }