feat(ece): support custom ownership type urns in ECE generation (data…

…hub-project#10999)
acryldata · Jul 26, 2024 · fdbcb68 · fdbcb68
1 parent bc75f7a
commit fdbcb68
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 6 deletions.
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/timeline/data/entity/OwnerChangeEvent.java b/metadata-io/src/main/java/com/linkedin/metadata/timeline/data/entity/OwnerChangeEvent.java
@@ -27,17 +27,29 @@ public OwnerChangeEvent(
       SemanticChangeType semVerChange,
       String description,
       Urn ownerUrn,
-      OwnershipType ownerType) {
+      OwnershipType ownerType,
+      Urn ownerTypeUrn) {
     super(
         entityUrn,
         category,
         operation,
         modifier,
-        ImmutableMap.of(
-            "ownerUrn", ownerUrn.toString(),
-            "ownerType", ownerType.toString()),
+        buildParameters(ownerUrn, ownerType, ownerTypeUrn),
         auditStamp,
         semVerChange,
         description);
   }
+
+  private static ImmutableMap<String, Object> buildParameters(
+      Urn ownerUrn, OwnershipType ownerType, Urn ownerTypeUrn) {
+    ImmutableMap.Builder<String, Object> builder =
+        new ImmutableMap.Builder<String, Object>()
+            .put("ownerUrn", ownerUrn.toString())
+            .put("ownerType", ownerType.toString());
+    if (ownerTypeUrn != null) {
+      builder.put("ownerTypeUrn", ownerTypeUrn.toString());
+    }
+
+    return builder.build();
+  }
 }
diff --git a/...ain/java/com/linkedin/metadata/timeline/eventgenerator/OwnershipChangeEventGenerator.java b/...ain/java/com/linkedin/metadata/timeline/eventgenerator/OwnershipChangeEventGenerator.java
@@ -62,6 +62,7 @@ private static List<ChangeEvent> computeDiffs(
                           entityUrn))
                   .ownerUrn(targetOwner.getOwner())
                   .ownerType(targetOwner.getType())
+                  .ownerTypeUrn(targetOwner.getTypeUrn())
                   .auditStamp(auditStamp)
                   .build());
         }
@@ -84,6 +85,7 @@ private static List<ChangeEvent> computeDiffs(
                         entityUrn))
                 .ownerUrn(baseOwner.getOwner())
                 .ownerType(baseOwner.getType())
+                .ownerTypeUrn(baseOwner.getTypeUrn())
                 .auditStamp(auditStamp)
                 .build());
         ++baseOwnerIdx;
@@ -104,6 +106,7 @@ private static List<ChangeEvent> computeDiffs(
                         entityUrn))
                 .ownerUrn(targetOwner.getOwner())
                 .ownerType(targetOwner.getType())
+                .ownerTypeUrn(targetOwner.getTypeUrn())
                 .auditStamp(auditStamp)
                 .build());
         ++targetOwnerIdx;
@@ -128,6 +131,7 @@ private static List<ChangeEvent> computeDiffs(
                       entityUrn))
               .ownerUrn(baseOwner.getOwner())
               .ownerType(baseOwner.getType())
+              .ownerTypeUrn(baseOwner.getTypeUrn())
               .auditStamp(auditStamp)
               .build());
       ++baseOwnerIdx;
@@ -150,6 +154,7 @@ private static List<ChangeEvent> computeDiffs(
                       entityUrn))
               .ownerUrn(targetOwner.getOwner())
               .ownerType(targetOwner.getType())
+              .ownerTypeUrn(targetOwner.getTypeUrn())
               .auditStamp(auditStamp)
               .build());
       ++targetOwnerIdx;

diff --git a/.../test/java/com/linkedin/metadata/kafka/hook/event/EntityChangeEventGeneratorHookTest.java b/.../test/java/com/linkedin/metadata/kafka/hook/event/EntityChangeEventGeneratorHookTest.java
@@ -309,11 +309,16 @@ public void testInvokeEntityOwnerChange() throws Exception {
     final Ownership newOwners = new Ownership();
     final Urn ownerUrn1 = Urn.createFromString("urn:li:corpuser:test1");
     final Urn ownerUrn2 = Urn.createFromString("urn:li:corpuser:test2");
+    final Urn ownerUrn3 = Urn.createFromString("urn:li:corpuser:test3");
     newOwners.setOwners(
         new OwnerArray(
             ImmutableList.of(
                 new Owner().setOwner(ownerUrn1).setType(OwnershipType.TECHNICAL_OWNER),
-                new Owner().setOwner(ownerUrn2).setType(OwnershipType.BUSINESS_OWNER))));
+                new Owner().setOwner(ownerUrn2).setType(OwnershipType.BUSINESS_OWNER),
+                new Owner()
+                    .setOwner(ownerUrn3)
+                    .setType(OwnershipType.CUSTOM)
+                    .setTypeUrn(Urn.createFromString("urn:li:ownershipType:my_custom_type")))));
     final Ownership prevOwners = new Ownership();
     prevOwners.setOwners(new OwnerArray());
     event.setAspect(GenericRecordUtils.serializeAspect(newOwners));
@@ -354,7 +359,24 @@ public void testInvokeEntityOwnerChange() throws Exception {
                 "ownerType",
                 OwnershipType.BUSINESS_OWNER.toString()),
             actorUrn);
-    verifyProducePlatformEvent(_mockClient, platformEvent2, true);
+    verifyProducePlatformEvent(_mockClient, platformEvent2, false);
+
+    PlatformEvent platformEvent3 =
+        createChangeEvent(
+            DATASET_ENTITY_NAME,
+            Urn.createFromString(TEST_DATASET_URN),
+            ChangeCategory.OWNER,
+            ChangeOperation.ADD,
+            ownerUrn3.toString(),
+            ImmutableMap.of(
+                "ownerUrn",
+                ownerUrn3.toString(),
+                "ownerType",
+                OwnershipType.CUSTOM.toString(),
+                "ownerTypeUrn",
+                "urn:li:ownershipType:my_custom_type"),
+            actorUrn);
+    verifyProducePlatformEvent(_mockClient, platformEvent3, true);
   }
 
   @Test