fix(gms/patch): Fix Upstream lineage patching when path contained enc…

…oded slash (datahub-project#11957)
acryldata · Nov 26, 2024 · 1ba1b2c · 1ba1b2c
1 parent 094433c
commit 1ba1b2c
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 3 deletions.
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
@@ -81,14 +81,13 @@ public static JsonNode populateTopLevelKeys(JsonNode transformedNode, JsonPatch
           PatchOperationType.REMOVE.equals(operationPath.getFirst())
               ? keys.length
               : keys.length - 1;
-
       // Skip first as it will always be blank due to path starting with /
       for (int i = 1; i < endIdx; i++) {
+        String decodedKey = decodeValue(keys[i]);
         if (parent.get(keys[i]) == null) {
-          String decodedKey = decodeValue(keys[i]);
           ((ObjectNode) parent).set(decodedKey, instance.objectNode());
         }
-        parent = parent.get(keys[i]);
+        parent = parent.get(decodedKey);
       }
     }
 

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
@@ -185,6 +185,29 @@ public void testPatchUpstream() throws Exception {
     // New entry in array because of new transformation type
     assertEquals(result4.getFineGrainedLineages().get(3), fineGrainedLineage4);
 
+    JsonPatchBuilder patchOperations5 = Json.createPatchBuilder();
+    String urn4 =
+        "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:s3,test-bucket/hive/folder_1/folder_2/my_dataset,DEV),c2)";
+    UrnArray downstreamUrns5 = new UrnArray();
+    downstreamUrns5.add(Urn.createFromString(urn4));
+    patchOperations5.add(
+        "/fineGrainedLineages/TRANSFORM/urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:s3,test-bucket~1hive~1folder_1~1folder_2~1my_dataset,DEV),c2)/urn:li:query:anotherQuery/urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:bigquery,upstream_table_1,PROD),c2)",
+        finegrainedLineageNode5.build());
+    JsonPatch jsonPatch5 = patchOperations5.build();
+    UpstreamLineage result5 = upstreamLineageTemplate.applyPatch(result4, jsonPatch5);
+    // Hack because Jackson parses values to doubles instead of floats
+    DataMap dataMap5 = new DataMap();
+    dataMap5.put("confidenceScore", 1.0);
+    FineGrainedLineage fineGrainedLineage5 = new FineGrainedLineage(dataMap5);
+    fineGrainedLineage5.setUpstreams(upstreamUrns3);
+    fineGrainedLineage5.setDownstreams(downstreamUrns5);
+    fineGrainedLineage5.setTransformOperation("TRANSFORM");
+    fineGrainedLineage5.setUpstreamType(FineGrainedLineageUpstreamType.FIELD_SET);
+    fineGrainedLineage5.setDownstreamType(FineGrainedLineageDownstreamType.FIELD);
+    fineGrainedLineage5.setQuery(UrnUtils.getUrn("urn:li:query:anotherQuery"));
+    // New entry in array because of new transformation type
+    assertEquals(result5.getFineGrainedLineages().get(4), fineGrainedLineage5);
+
     // Remove
     JsonPatchBuilder removeOperations = Json.createPatchBuilder();
     removeOperations.remove(