fix(forms) Consolidate how we add owners in forms and elsewhere (#14449)

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/OwnerUtils.java

@@ -5,10 +5,8 @@ import static com.linkedin.datahub.graphql.resolvers.mutate.MutationUtils.*;
 import com.datahub.authorization.ConjunctivePrivilegeGroup;
 import com.datahub.authorization.DisjunctivePrivilegeGroup;
 import com.google.common.collect.ImmutableList;
-import com.linkedin.common.Owner;
 import com.linkedin.common.OwnerArray;
 import com.linkedin.common.Ownership;
-import com.linkedin.common.OwnershipSource;
 import com.linkedin.common.OwnershipSourceType;
 import com.linkedin.common.urn.CorpuserUrn;
 import com.linkedin.common.urn.Urn;
@@ -25,6 +23,7 @@ import com.linkedin.metadata.Constants;
 import com.linkedin.metadata.authorization.PoliciesConfig;
 import com.linkedin.metadata.entity.EntityService;
 import com.linkedin.metadata.entity.EntityUtils;
+import com.linkedin.metadata.service.util.OwnerServiceUtils;
 import com.linkedin.mxe.MetadataChangeProposal;
 import io.datahubproject.metadata.context.OperationContext;
 import java.util.ArrayList;
@@ -100,11 +99,16 @@ public class OwnerUtils {
                 new Ownership());
     ownershipAspect.setLastModified(EntityUtils.getAuditStamp(actor));
     for (OwnerInput input : owners) {
-      addOwnerToAspect(
+      final OwnershipType ownershipType =
+          input.getType() != null
+              ? OwnershipType.valueOf(input.getType().toString())
+              : OwnershipType.NONE;
+      OwnerServiceUtils.addOwnerToAspect(
           ownershipAspect,
           UrnUtils.getUrn(input.getOwnerUrn()),
-          input.getType(),
+          com.linkedin.common.OwnershipType.valueOf(ownershipType.toString()),
-          UrnUtils.getUrn(input.getOwnershipTypeUrn()));
+          UrnUtils.getUrn(input.getOwnershipTypeUrn()),
+          OwnershipSourceType.MANUAL);
     }
     return buildMetadataChangeProposalWithUrn(
         resourceUrn, Constants.OWNERSHIP_ASPECT_NAME, ownershipAspect);
@@ -131,58 +135,6 @@ public class OwnerUtils {
         resourceUrn, Constants.OWNERSHIP_ASPECT_NAME, ownershipAspect);
   }
 
-  private static void addOwnerToAspect(
-      Ownership ownershipAspect, Urn ownerUrn, OwnershipType type, Urn ownershipTypeUrn) {
-    if (!ownershipAspect.hasOwners()) {
-      ownershipAspect.setOwners(new OwnerArray());
-    }
-
-    OwnerArray ownerArray = new OwnerArray(ownershipAspect.getOwners());
-    removeExistingOwnerIfExists(ownerArray, ownerUrn, ownershipTypeUrn);
-
-    Owner newOwner = new Owner();
-
-    // For backwards compatibility we have to always set the deprecated type.
-    // If the type exists we assume it's an old ownership type that we can map to.
-    // Else if it's a net new custom ownership type set old type to CUSTOM.
-    com.linkedin.common.OwnershipType gmsType =
-        type != null
-            ? com.linkedin.common.OwnershipType.valueOf(type.toString())
-            : com.linkedin.common.OwnershipType.CUSTOM;
-
-    newOwner.setType(gmsType);
-    newOwner.setTypeUrn(ownershipTypeUrn);
-    newOwner.setSource(new OwnershipSource().setType(OwnershipSourceType.MANUAL));
-    newOwner.setOwner(ownerUrn);
-    ownerArray.add(newOwner);
-    ownershipAspect.setOwners(ownerArray);
-  }
-
-  private static void removeExistingOwnerIfExists(
-      OwnerArray ownerArray, Urn ownerUrn, Urn ownershipTypeUrn) {
-    ownerArray.removeIf(
-        owner -> {
-          // Remove old ownership if it exists (check ownerUrn + type (entity & deprecated type))
-          return isOwnerEqual(owner, ownerUrn, ownershipTypeUrn);
-        });
-  }
-
-  public static boolean isOwnerEqual(
-      @Nonnull Owner owner, @Nonnull Urn ownerUrn, @Nullable Urn ownershipTypeUrn) {
-    if (!owner.getOwner().equals(ownerUrn)) {
-      return false;
-    }
-    if (owner.getTypeUrn() != null && ownershipTypeUrn != null) {
-      return owner.getTypeUrn().equals(ownershipTypeUrn);
-    }
-    if (ownershipTypeUrn == null) {
-      return true;
-    }
-    // Fall back to mapping deprecated type to the new ownership entity
-    return mapOwnershipTypeToEntity(OwnershipType.valueOf(owner.getType().toString()).name())
-        .equals(ownershipTypeUrn.toString());
-  }
-
   private static void removeOwnersIfExists(
       Ownership ownershipAspect, List<Urn> ownerUrns, Urn ownershipTypeUrn) {
     if (!ownershipAspect.hasOwners()) {
@@ -191,7 +143,7 @@ public class OwnerUtils {
 
     OwnerArray ownerArray = ownershipAspect.getOwners();
     for (Urn ownerUrn : ownerUrns) {
-      removeExistingOwnerIfExists(ownerArray, ownerUrn, ownershipTypeUrn);
+      OwnerServiceUtils.removeExistingOwnerIfExists(ownerArray, ownerUrn, ownershipTypeUrn);
     }
   }
 

datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/utils/OwnerUtilsTest.java

@@ -1,116 +0,0 @@
-package com.linkedin.datahub.graphql.utils;
-
-import static org.testng.AssertJUnit.*;
-
-import com.linkedin.common.Owner;
-import com.linkedin.common.OwnershipType;
-import com.linkedin.common.urn.Urn;
-import com.linkedin.datahub.graphql.resolvers.mutate.util.OwnerUtils;
-import java.net.URISyntaxException;
-import org.testng.annotations.Test;
-
-public class OwnerUtilsTest {
-
-  public static String TECHNICAL_OWNER_OWNERSHIP_TYPE_URN =
-      "urn:li:ownershipType:__system__technical_owner";
-  public static String BUSINESS_OWNER_OWNERSHIP_TYPE_URN =
-      "urn:li:ownershipType:__system__business_owner";
-
-  @Test
-  public void testMapOwnershipType() {
-    assertEquals(
-        OwnerUtils.mapOwnershipTypeToEntity("TECHNICAL_OWNER"), TECHNICAL_OWNER_OWNERSHIP_TYPE_URN);
-  }
-
-  @Test
-  public void testIsOwnerEqualUrnOnly() throws URISyntaxException {
-    Urn ownerUrn1 = new Urn("urn:li:corpuser:foo");
-    Owner owner1 = new Owner();
-    owner1.setOwner(ownerUrn1);
-    assertTrue(OwnerUtils.isOwnerEqual(owner1, ownerUrn1, null));
-
-    Urn ownerUrn2 = new Urn("urn:li:corpuser:bar");
-    assertFalse(OwnerUtils.isOwnerEqual(owner1, ownerUrn2, null));
-  }
-
-  @Test
-  public void testIsOwnerEqualWithLegacyTypeOnly() throws URISyntaxException {
-
-    Urn technicalOwnershipTypeUrn = new Urn(TECHNICAL_OWNER_OWNERSHIP_TYPE_URN);
-    Urn ownerUrn1 = new Urn("urn:li:corpuser:foo");
-    Owner ownerWithTechnicalOwnership = new Owner();
-    ownerWithTechnicalOwnership.setOwner(ownerUrn1);
-    ownerWithTechnicalOwnership.setType(OwnershipType.TECHNICAL_OWNER);
-
-    assertTrue(
-        OwnerUtils.isOwnerEqual(ownerWithTechnicalOwnership, ownerUrn1, technicalOwnershipTypeUrn));
-
-    Owner ownerWithBusinessOwnership = new Owner();
-    ownerWithBusinessOwnership.setOwner(ownerUrn1);
-    ownerWithBusinessOwnership.setType(OwnershipType.BUSINESS_OWNER);
-    assertFalse(
-        OwnerUtils.isOwnerEqual(
-            ownerWithBusinessOwnership, ownerUrn1, new Urn(TECHNICAL_OWNER_OWNERSHIP_TYPE_URN)));
-  }
-
-  @Test
-  public void testIsOwnerEqualOnlyOwnershipTypeUrn() throws URISyntaxException {
-
-    Urn technicalOwnershipTypeUrn = new Urn(TECHNICAL_OWNER_OWNERSHIP_TYPE_URN);
-    Urn businessOwnershipTypeUrn = new Urn(BUSINESS_OWNER_OWNERSHIP_TYPE_URN);
-    Urn ownerUrn1 = new Urn("urn:li:corpuser:foo");
-    Urn ownerUrn2 = new Urn("urn:li:corpuser:bar");
-
-    Owner ownerWithTechnicalOwnership = new Owner();
-    ownerWithTechnicalOwnership.setOwner(ownerUrn1);
-    ownerWithTechnicalOwnership.setTypeUrn(technicalOwnershipTypeUrn);
-
-    Owner ownerWithBusinessOwnership = new Owner();
-    ownerWithBusinessOwnership.setOwner(ownerUrn1);
-    ownerWithBusinessOwnership.setTypeUrn(businessOwnershipTypeUrn);
-
-    Owner ownerWithoutOwnershipType = new Owner();
-    ownerWithoutOwnershipType.setOwner(ownerUrn1);
-    ownerWithoutOwnershipType.setType(OwnershipType.NONE);
-
-    Owner owner2WithoutOwnershipType = new Owner();
-    owner2WithoutOwnershipType.setOwner(ownerUrn2);
-    owner2WithoutOwnershipType.setType(OwnershipType.NONE);
-
-    assertTrue(
-        OwnerUtils.isOwnerEqual(ownerWithTechnicalOwnership, ownerUrn1, technicalOwnershipTypeUrn));
-    assertFalse(
-        OwnerUtils.isOwnerEqual(ownerWithBusinessOwnership, ownerUrn1, technicalOwnershipTypeUrn));
-    assertTrue(OwnerUtils.isOwnerEqual(ownerWithTechnicalOwnership, ownerUrn1, null));
-    assertTrue(OwnerUtils.isOwnerEqual(ownerWithoutOwnershipType, ownerUrn1, null));
-    assertFalse(OwnerUtils.isOwnerEqual(owner2WithoutOwnershipType, ownerUrn1, null));
-  }
-
-  public void testIsOwnerEqualWithBothLegacyAndNewType() throws URISyntaxException {
-    Urn technicalOwnershipTypeUrn = new Urn(TECHNICAL_OWNER_OWNERSHIP_TYPE_URN);
-    Urn businessOwnershipTypeUrn = new Urn(BUSINESS_OWNER_OWNERSHIP_TYPE_URN);
-    Urn ownerUrn1 = new Urn("urn:li:corpuser:foo");
-
-    Owner ownerWithLegacyTechnicalOwnership = new Owner();
-    ownerWithLegacyTechnicalOwnership.setOwner(ownerUrn1);
-    ownerWithLegacyTechnicalOwnership.setType(OwnershipType.TECHNICAL_OWNER);
-
-    assertTrue(
-        OwnerUtils.isOwnerEqual(
-            ownerWithLegacyTechnicalOwnership, ownerUrn1, technicalOwnershipTypeUrn));
-    assertFalse(
-        OwnerUtils.isOwnerEqual(
-            ownerWithLegacyTechnicalOwnership, ownerUrn1, businessOwnershipTypeUrn));
-
-    Owner ownerWithNewTechnicalOwnership = new Owner();
-    ownerWithLegacyTechnicalOwnership.setOwner(ownerUrn1);
-    ownerWithLegacyTechnicalOwnership.setTypeUrn(technicalOwnershipTypeUrn);
-
-    assertTrue(
-        OwnerUtils.isOwnerEqual(
-            ownerWithNewTechnicalOwnership, ownerUrn1, technicalOwnershipTypeUrn));
-    assertFalse(
-        OwnerUtils.isOwnerEqual(
-            ownerWithNewTechnicalOwnership, ownerUrn1, businessOwnershipTypeUrn));
-  }
-}

metadata-io/src/test/java/com/linkedin/metadata/service/OwnerServiceUtilsTest.java

@@ -0,0 +1,485 @@
+package com.linkedin.metadata.service;
+
+import static org.testng.Assert.*;
+
+import com.linkedin.common.Owner;
+import com.linkedin.common.OwnerArray;
+import com.linkedin.common.Ownership;
+import com.linkedin.common.OwnershipSource;
+import com.linkedin.common.OwnershipSourceType;
+import com.linkedin.common.OwnershipType;
+import com.linkedin.common.urn.Urn;
+import com.linkedin.common.urn.UrnUtils;
+import com.linkedin.metadata.service.util.OwnerServiceUtils;
+import java.net.URISyntaxException;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+public class OwnerServiceUtilsTest {
+
+  public static String TECHNICAL_OWNER_OWNERSHIP_TYPE_URN =
+      "urn:li:ownershipType:__system__technical_owner";
+  public static String BUSINESS_OWNER_OWNERSHIP_TYPE_URN =
+      "urn:li:ownershipType:__system__business_owner";
+  public static String DATA_STEWARD_OWNERSHIP_TYPE_URN =
+      "urn:li:ownershipType:__system__data_steward";
+
+  private Urn ownerUrn1;
+  private Urn ownerUrn2;
+  private Urn technicalOwnershipTypeUrn;
+  private Urn businessOwnershipTypeUrn;
+  private Urn dataStewardOwnershipTypeUrn;
+
+  @BeforeMethod
+  public void setUp() throws URISyntaxException {
+    ownerUrn1 = new Urn("urn:li:corpuser:foo");
+    ownerUrn2 = new Urn("urn:li:corpuser:bar");
+    technicalOwnershipTypeUrn = new Urn(TECHNICAL_OWNER_OWNERSHIP_TYPE_URN);
+    businessOwnershipTypeUrn = new Urn(BUSINESS_OWNER_OWNERSHIP_TYPE_URN);
+    dataStewardOwnershipTypeUrn = new Urn(DATA_STEWARD_OWNERSHIP_TYPE_URN);
+  }
+
+  // Tests for addOwnerToAspect method
+
+  @Test
+  public void testAddOwnerToAspectWithoutExistingOwners() {
+    Ownership ownership = new Ownership();
+
+    OwnerServiceUtils.addOwnerToAspect(
+        ownership,
+        ownerUrn1,
+        OwnershipType.TECHNICAL_OWNER,
+        technicalOwnershipTypeUrn,
+        OwnershipSourceType.MANUAL);
+
+    assertTrue(ownership.hasOwners());
+    assertEquals(1, ownership.getOwners().size());
+
+    Owner addedOwner = ownership.getOwners().get(0);
+    assertEquals(ownerUrn1, addedOwner.getOwner());
+    assertEquals(OwnershipType.TECHNICAL_OWNER, addedOwner.getType());
+    assertEquals(technicalOwnershipTypeUrn, addedOwner.getTypeUrn());
+    assertEquals(OwnershipSourceType.MANUAL, addedOwner.getSource().getType());
+  }
+
+  @Test
+  public void testAddOwnerToAspectWithExistingOwners() {
+    Ownership ownership = new Ownership();
+    OwnerArray existingOwners = new OwnerArray();
+
+    // Add an existing owner
+    Owner existingOwner = new Owner();
+    existingOwner.setOwner(ownerUrn2);
+    existingOwner.setType(OwnershipType.BUSINESS_OWNER);
+    existingOwner.setTypeUrn(businessOwnershipTypeUrn);
+    existingOwner.setSource(new OwnershipSource().setType(OwnershipSourceType.MANUAL));
+    existingOwners.add(existingOwner);
+
+    ownership.setOwners(existingOwners);
+
+    OwnerServiceUtils.addOwnerToAspect(
+        ownership, ownerUrn1, OwnershipType.TECHNICAL_OWNER, technicalOwnershipTypeUrn);
+
+    assertEquals(2, ownership.getOwners().size());
+
+    // Verify the new owner was added
+    Owner newOwner = ownership.getOwners().get(1);
+    assertEquals(ownerUrn1, newOwner.getOwner());
+    assertEquals(OwnershipType.TECHNICAL_OWNER, newOwner.getType());
+    assertEquals(technicalOwnershipTypeUrn, newOwner.getTypeUrn());
+  }
+
+  @Test
+  public void testAddOwnerToAspectReplacesExistingOwnerSameTypeUrn() {
+    Ownership ownership = new Ownership();
+    OwnerArray existingOwners = new OwnerArray();
+
+    // Add an existing owner with same URN and type
+    Owner existingOwner = new Owner();
+    existingOwner.setOwner(ownerUrn1);
+    existingOwner.setType(OwnershipType.TECHNICAL_OWNER);
+    existingOwner.setTypeUrn(technicalOwnershipTypeUrn);
+    existingOwner.setSource(new OwnershipSource().setType(OwnershipSourceType.MANUAL));
+    existingOwners.add(existingOwner);
+
+    ownership.setOwners(existingOwners);
+
+    OwnerServiceUtils.addOwnerToAspect(
+        ownership, ownerUrn1, OwnershipType.TECHNICAL_OWNER, technicalOwnershipTypeUrn);
+
+    // Should still only have 1 owner (the existing one was replaced)
+    assertEquals(1, ownership.getOwners().size());
+
+    Owner updatedOwner = ownership.getOwners().get(0);
+    assertEquals(ownerUrn1, updatedOwner.getOwner());
+    assertEquals(OwnershipType.TECHNICAL_OWNER, updatedOwner.getType());
+    assertEquals(technicalOwnershipTypeUrn, updatedOwner.getTypeUrn());
+  }
+
+  @Test
+  public void testAddOwnerToAspectWithDifferentOwnershipTypes() {
+    Ownership ownership = new Ownership();
+
+    // Add technical owner
+    OwnerServiceUtils.addOwnerToAspect(
+        ownership, ownerUrn1, OwnershipType.TECHNICAL_OWNER, technicalOwnershipTypeUrn);
+
+    // Add business owner for same user
+    OwnerServiceUtils.addOwnerToAspect(
+        ownership, ownerUrn1, OwnershipType.BUSINESS_OWNER, businessOwnershipTypeUrn);
+
+    assertEquals(2, ownership.getOwners().size());
+
+    // Both ownership types should be present for the same user
+    assertTrue(
+        ownership.getOwners().stream()
+            .anyMatch(
+                owner ->
+                    owner.getOwner().equals(ownerUrn1)
+                        && owner.getTypeUrn().equals(technicalOwnershipTypeUrn)));
+    assertTrue(
+        ownership.getOwners().stream()
+            .anyMatch(
+                owner ->
+                    owner.getOwner().equals(ownerUrn1)
+                        && owner.getTypeUrn().equals(businessOwnershipTypeUrn)));
+  }
+
+  // Tests for removeExistingOwnerIfExists method
+
+  @Test
+  public void testRemoveExistingOwnerIfExistsWithMatchingOwner() {
+    OwnerArray ownerArray = new OwnerArray();
+
+    // Add multiple owners
+    Owner owner1 = new Owner();
+    owner1.setOwner(ownerUrn1);
+    owner1.setType(OwnershipType.TECHNICAL_OWNER);
+    owner1.setTypeUrn(technicalOwnershipTypeUrn);
+    ownerArray.add(owner1);
+
+    Owner owner2 = new Owner();
+    owner2.setOwner(ownerUrn2);
+    owner2.setType(OwnershipType.BUSINESS_OWNER);
+    owner2.setTypeUrn(businessOwnershipTypeUrn);
+    ownerArray.add(owner2);
+
+    assertEquals(2, ownerArray.size());
+
+    // Remove owner1
+    OwnerServiceUtils.removeExistingOwnerIfExists(ownerArray, ownerUrn1, technicalOwnershipTypeUrn);
+
+    assertEquals(1, ownerArray.size());
+    assertEquals(ownerUrn2, ownerArray.get(0).getOwner());
+  }
+
+  @Test
+  public void testRemoveExistingOwnerIfExistsWithNoMatch() {
+    OwnerArray ownerArray = new OwnerArray();
+
+    Owner owner1 = new Owner();
+    owner1.setOwner(ownerUrn1);
+    owner1.setType(OwnershipType.TECHNICAL_OWNER);
+    owner1.setTypeUrn(technicalOwnershipTypeUrn);
+    ownerArray.add(owner1);
+
+    assertEquals(1, ownerArray.size());
+
+    // Try to remove non-existent owner
+    OwnerServiceUtils.removeExistingOwnerIfExists(ownerArray, ownerUrn2, businessOwnershipTypeUrn);
+
+    // Should still have 1 owner
+    assertEquals(1, ownerArray.size());
+    assertEquals(ownerUrn1, ownerArray.get(0).getOwner());
+  }
+
+  @Test
+  public void testRemoveExistingOwnerIfExistsWithEmptyArray() {
+    OwnerArray ownerArray = new OwnerArray();
+
+    assertEquals(0, ownerArray.size());
+
+    // Try to remove from empty array
+    OwnerServiceUtils.removeExistingOwnerIfExists(ownerArray, ownerUrn1, technicalOwnershipTypeUrn);
+
+    // Should still be empty
+    assertEquals(0, ownerArray.size());
+  }
+
+  @Test
+  public void testRemoveExistingOwnerIfExistsWithNullTypeUrn() {
+    OwnerArray ownerArray = new OwnerArray();
+
+    Owner owner1 = new Owner();
+    owner1.setOwner(ownerUrn1);
+    owner1.setType(OwnershipType.TECHNICAL_OWNER);
+    ownerArray.add(owner1);
+
+    assertEquals(1, ownerArray.size());
+
+    // Remove with null type URN - should match any owner with matching URN
+    OwnerServiceUtils.removeExistingOwnerIfExists(ownerArray, ownerUrn1, null);
+
+    assertEquals(0, ownerArray.size());
+  }
+
+  // Tests for mapOwnershipTypeToEntity method
+
+  @Test
+  public void testMapOwnershipType() {
+    assertEquals(
+        OwnerServiceUtils.mapOwnershipTypeToEntity("TECHNICAL_OWNER"),
+        UrnUtils.getUrn(TECHNICAL_OWNER_OWNERSHIP_TYPE_URN));
+  }
+
+  @Test
+  public void testMapOwnershipTypeBusinessOwner() {
+    assertEquals(
+        OwnerServiceUtils.mapOwnershipTypeToEntity("BUSINESS_OWNER"),
+        UrnUtils.getUrn(BUSINESS_OWNER_OWNERSHIP_TYPE_URN));
+  }
+
+  @Test
+  public void testMapOwnershipTypeDataSteward() {
+    assertEquals(
+        OwnerServiceUtils.mapOwnershipTypeToEntity("DATA_STEWARD"),
+        UrnUtils.getUrn(DATA_STEWARD_OWNERSHIP_TYPE_URN));
+  }
+
+  @Test
+  public void testMapOwnershipTypeLowercase() {
+    assertEquals(
+        OwnerServiceUtils.mapOwnershipTypeToEntity("technical_owner"),
+        UrnUtils.getUrn(TECHNICAL_OWNER_OWNERSHIP_TYPE_URN));
+  }
+
+  @Test
+  public void testMapOwnershipTypeMixedCase() {
+    assertEquals(
+        OwnerServiceUtils.mapOwnershipTypeToEntity("Technical_Owner"),
+        UrnUtils.getUrn(TECHNICAL_OWNER_OWNERSHIP_TYPE_URN));
+  }
+
+  @Test
+  public void testMapOwnershipTypeWithSystemIdConstant() {
+    String result = OwnerServiceUtils.mapOwnershipTypeToEntity("CUSTOM_TYPE").toString();
+    assertTrue(result.contains(OwnerServiceUtils.SYSTEM_ID));
+    assertTrue(result.contains("custom_type"));
+  }
+
+  @Test
+  public void testMapOwnershipTypeEmptyString() {
+    Urn result = OwnerServiceUtils.mapOwnershipTypeToEntity("");
+    String expectedUrn = "urn:li:ownershipType:" + OwnerServiceUtils.SYSTEM_ID;
+    assertEquals(expectedUrn, result.toString());
+  }
+
+  @Test
+  public void testMapOwnershipTypeSpecialCharacters() {
+    // Test with special characters that get converted to lowercase
+    Urn result = OwnerServiceUtils.mapOwnershipTypeToEntity("SPECIAL-TYPE_123");
+    String expectedUrn = "urn:li:ownershipType:" + OwnerServiceUtils.SYSTEM_ID + "special-type_123";
+    assertEquals(expectedUrn, result.toString());
+  }
+
+  @Test
+  public void testIsOwnerEqualUrnOnly() throws URISyntaxException {
+    Urn ownerUrn1 = new Urn("urn:li:corpuser:foo");
+    Owner owner1 = new Owner();
+    owner1.setOwner(ownerUrn1);
+    assertTrue(OwnerServiceUtils.isOwnerEqual(owner1, ownerUrn1, null));
+
+    Urn ownerUrn2 = new Urn("urn:li:corpuser:bar");
+    assertFalse(OwnerServiceUtils.isOwnerEqual(owner1, ownerUrn2, null));
+  }
+
+  @Test
+  public void testIsOwnerEqualWithLegacyTypeOnly() throws URISyntaxException {
+
+    Urn technicalOwnershipTypeUrn = new Urn(TECHNICAL_OWNER_OWNERSHIP_TYPE_URN);
+    Urn ownerUrn1 = new Urn("urn:li:corpuser:foo");
+    Owner ownerWithTechnicalOwnership = new Owner();
+    ownerWithTechnicalOwnership.setOwner(ownerUrn1);
+    ownerWithTechnicalOwnership.setType(OwnershipType.TECHNICAL_OWNER);
+
+    assertTrue(
+        OwnerServiceUtils.isOwnerEqual(
+            ownerWithTechnicalOwnership, ownerUrn1, technicalOwnershipTypeUrn));
+
+    Owner ownerWithBusinessOwnership = new Owner();
+    ownerWithBusinessOwnership.setOwner(ownerUrn1);
+    ownerWithBusinessOwnership.setType(OwnershipType.BUSINESS_OWNER);
+    assertFalse(
+        OwnerServiceUtils.isOwnerEqual(
+            ownerWithBusinessOwnership, ownerUrn1, new Urn(TECHNICAL_OWNER_OWNERSHIP_TYPE_URN)));
+  }
+
+  @Test
+  public void testIsOwnerEqualOnlyOwnershipTypeUrn() throws URISyntaxException {
+
+    Urn technicalOwnershipTypeUrn = new Urn(TECHNICAL_OWNER_OWNERSHIP_TYPE_URN);
+    Urn businessOwnershipTypeUrn = new Urn(BUSINESS_OWNER_OWNERSHIP_TYPE_URN);
+    Urn ownerUrn1 = new Urn("urn:li:corpuser:foo");
+    Urn ownerUrn2 = new Urn("urn:li:corpuser:bar");
+
+    Owner ownerWithTechnicalOwnership = new Owner();
+    ownerWithTechnicalOwnership.setOwner(ownerUrn1);
+    ownerWithTechnicalOwnership.setTypeUrn(technicalOwnershipTypeUrn);
+
+    Owner ownerWithBusinessOwnership = new Owner();
+    ownerWithBusinessOwnership.setOwner(ownerUrn1);
+    ownerWithBusinessOwnership.setTypeUrn(businessOwnershipTypeUrn);
+
+    Owner ownerWithoutOwnershipType = new Owner();
+    ownerWithoutOwnershipType.setOwner(ownerUrn1);
+    ownerWithoutOwnershipType.setType(OwnershipType.NONE);
+
+    Owner owner2WithoutOwnershipType = new Owner();
+    owner2WithoutOwnershipType.setOwner(ownerUrn2);
+    owner2WithoutOwnershipType.setType(OwnershipType.NONE);
+
+    assertTrue(
+        OwnerServiceUtils.isOwnerEqual(
+            ownerWithTechnicalOwnership, ownerUrn1, technicalOwnershipTypeUrn));
+    assertFalse(
+        OwnerServiceUtils.isOwnerEqual(
+            ownerWithBusinessOwnership, ownerUrn1, technicalOwnershipTypeUrn));
+    assertTrue(OwnerServiceUtils.isOwnerEqual(ownerWithTechnicalOwnership, ownerUrn1, null));
+    assertTrue(OwnerServiceUtils.isOwnerEqual(ownerWithoutOwnershipType, ownerUrn1, null));
+    assertFalse(OwnerServiceUtils.isOwnerEqual(owner2WithoutOwnershipType, ownerUrn1, null));
+  }
+
+  @Test
+  public void testIsOwnerEqualWithBothLegacyAndNewType() throws URISyntaxException {
+    Urn technicalOwnershipTypeUrn = new Urn(TECHNICAL_OWNER_OWNERSHIP_TYPE_URN);
+    Urn businessOwnershipTypeUrn = new Urn(BUSINESS_OWNER_OWNERSHIP_TYPE_URN);
+    Urn ownerUrn1 = new Urn("urn:li:corpuser:foo");
+
+    Owner ownerWithLegacyTechnicalOwnership = new Owner();
+    ownerWithLegacyTechnicalOwnership.setOwner(ownerUrn1);
+    ownerWithLegacyTechnicalOwnership.setType(OwnershipType.TECHNICAL_OWNER);
+
+    assertTrue(
+        OwnerServiceUtils.isOwnerEqual(
+            ownerWithLegacyTechnicalOwnership, ownerUrn1, technicalOwnershipTypeUrn));
+    assertFalse(
+        OwnerServiceUtils.isOwnerEqual(
+            ownerWithLegacyTechnicalOwnership, ownerUrn1, businessOwnershipTypeUrn));
+
+    Owner ownerWithNewTechnicalOwnership = new Owner();
+    ownerWithNewTechnicalOwnership.setOwner(ownerUrn1);
+    ownerWithNewTechnicalOwnership.setTypeUrn(technicalOwnershipTypeUrn);
+
+    assertTrue(
+        OwnerServiceUtils.isOwnerEqual(
+            ownerWithNewTechnicalOwnership, ownerUrn1, technicalOwnershipTypeUrn));
+    assertFalse(
+        OwnerServiceUtils.isOwnerEqual(
+            ownerWithNewTechnicalOwnership, ownerUrn1, businessOwnershipTypeUrn));
+  }
+
+  // Additional edge case tests for isOwnerEqual method
+
+  @Test
+  public void testIsOwnerEqualWithNullOwnerTypeUrn() {
+    Owner owner = new Owner();
+    owner.setOwner(ownerUrn1);
+    owner.setType(OwnershipType.TECHNICAL_OWNER);
+    // No typeUrn set
+
+    assertTrue(OwnerServiceUtils.isOwnerEqual(owner, ownerUrn1, null));
+    assertTrue(OwnerServiceUtils.isOwnerEqual(owner, ownerUrn1, technicalOwnershipTypeUrn));
+  }
+
+  @Test
+  public void testIsOwnerEqualWithMismatchedUrns() {
+    Owner owner = new Owner();
+    owner.setOwner(ownerUrn1);
+    owner.setType(OwnershipType.TECHNICAL_OWNER);
+    owner.setTypeUrn(technicalOwnershipTypeUrn);
+
+    assertFalse(OwnerServiceUtils.isOwnerEqual(owner, ownerUrn2, technicalOwnershipTypeUrn));
+    assertFalse(OwnerServiceUtils.isOwnerEqual(owner, ownerUrn2, null));
+  }
+
+  @Test
+  public void testIsOwnerEqualWithLegacyTypeMapping() {
+    Owner owner = new Owner();
+    owner.setOwner(ownerUrn1);
+    owner.setType(OwnershipType.BUSINESS_OWNER);
+    // No typeUrn set, should fall back to legacy type mapping
+
+    assertTrue(OwnerServiceUtils.isOwnerEqual(owner, ownerUrn1, businessOwnershipTypeUrn));
+    assertFalse(OwnerServiceUtils.isOwnerEqual(owner, ownerUrn1, technicalOwnershipTypeUrn));
+  }
+
+  @Test
+  public void testIsOwnerEqualWithDataStewardType() {
+    Owner owner = new Owner();
+    owner.setOwner(ownerUrn1);
+    owner.setType(OwnershipType.DATA_STEWARD);
+
+    assertTrue(OwnerServiceUtils.isOwnerEqual(owner, ownerUrn1, dataStewardOwnershipTypeUrn));
+    assertFalse(OwnerServiceUtils.isOwnerEqual(owner, ownerUrn1, technicalOwnershipTypeUrn));
+  }
+
+  @Test
+  public void testIsOwnerEqualWithNoneType() {
+    Owner owner = new Owner();
+    owner.setOwner(ownerUrn1);
+    owner.setType(OwnershipType.NONE);
+
+    assertTrue(OwnerServiceUtils.isOwnerEqual(owner, ownerUrn1, null));
+    // When ownershipTypeUrn is provided but owner has NONE type, it should fall back to mapping
+    Urn noneTypeUrn = OwnerServiceUtils.mapOwnershipTypeToEntity("NONE");
+    assertTrue(OwnerServiceUtils.isOwnerEqual(owner, ownerUrn1, noneTypeUrn));
+  }
+
+  // Tests for SYSTEM_ID constant usage
+
+  @Test
+  public void testSystemIdConstant() {
+    assertEquals("__system__", OwnerServiceUtils.SYSTEM_ID);
+  }
+
+  @Test
+  public void testSystemIdUsedInMapping() {
+    String customType = "CUSTOM_OWNERSHIP";
+    Urn result = OwnerServiceUtils.mapOwnershipTypeToEntity(customType);
+    String expected =
+        "urn:li:ownershipType:" + OwnerServiceUtils.SYSTEM_ID + customType.toLowerCase();
+    assertEquals(expected, result.toString());
+  }
+
+  @Test
+  public void testSystemIdInAllStandardTypes() {
+    String[] standardTypes = {"TECHNICAL_OWNER", "BUSINESS_OWNER", "DATA_STEWARD"};
+
+    for (String type : standardTypes) {
+      Urn result = OwnerServiceUtils.mapOwnershipTypeToEntity(type);
+      assertTrue(result.toString().contains(OwnerServiceUtils.SYSTEM_ID));
+    }
+  }
+
+  @Test
+  public void testFullOwnershipWorkflowWithTypeUrns() {
+    Ownership ownership = new Ownership();
+
+    // Add owner using type URN
+    OwnerServiceUtils.addOwnerToAspect(
+        ownership, ownerUrn1, OwnershipType.TECHNICAL_OWNER, technicalOwnershipTypeUrn);
+
+    assertEquals(1, ownership.getOwners().size());
+    Owner addedOwner = ownership.getOwners().get(0);
+    assertEquals(ownerUrn1, addedOwner.getOwner());
+    assertEquals(technicalOwnershipTypeUrn, addedOwner.getTypeUrn());
+
+    // Verify exact match
+    assertTrue(OwnerServiceUtils.isOwnerEqual(addedOwner, ownerUrn1, technicalOwnershipTypeUrn));
+
+    // Verify removal works
+    OwnerServiceUtils.removeExistingOwnerIfExists(
+        ownership.getOwners(), ownerUrn1, technicalOwnershipTypeUrn);
+    assertEquals(0, ownership.getOwners().size());
+  }
+}

metadata-service/services/src/main/java/com/linkedin/metadata/service/OwnerService.java

@@ -4,7 +4,6 @@ import static com.linkedin.metadata.entity.AspectUtils.*;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.linkedin.common.AuditStamp;
-import com.linkedin.common.Owner;
 import com.linkedin.common.OwnerArray;
 import com.linkedin.common.Ownership;
 import com.linkedin.common.OwnershipType;
@@ -13,6 +12,7 @@ import com.linkedin.common.urn.UrnUtils;
 import com.linkedin.entity.client.SystemEntityClient;
 import com.linkedin.metadata.Constants;
 import com.linkedin.metadata.resource.ResourceReference;
+import com.linkedin.metadata.service.util.OwnerServiceUtils;
 import com.linkedin.mxe.MetadataChangeProposal;
 import io.datahubproject.metadata.context.OperationContext;
 import java.util.ArrayList;
@@ -20,13 +20,12 @@ import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 import lombok.extern.slf4j.Slf4j;
 
 @Slf4j
 public class OwnerService extends BaseService {
 
-  public static final String SYSTEM_ID = "__system__";
-
   public OwnerService(@Nonnull SystemEntityClient entityClient) {
     super(entityClient);
   }
@@ -42,10 +41,11 @@ public class OwnerService extends BaseService {
       @Nonnull OperationContext opContext,
       @Nonnull List<Urn> ownerUrns,
       @Nonnull List<ResourceReference> resources,
-      @Nonnull OwnershipType ownershipType) {
+      @Nonnull OwnershipType ownershipType,
+      @Nullable Urn ownershipTypeUrn) {
     log.debug("Batch adding Owners to entities. owners: {}, resources: {}", resources, ownerUrns);
     try {
-      addOwnersToResources(opContext, ownerUrns, resources, ownershipType);
+      addOwnersToResources(opContext, ownerUrns, resources, ownershipType, ownershipTypeUrn);
     } catch (Exception e) {
       throw new RuntimeException(
           String.format(
@@ -83,10 +83,11 @@ public class OwnerService extends BaseService {
       @Nonnull OperationContext opContext,
       List<com.linkedin.common.urn.Urn> ownerUrns,
       List<ResourceReference> resources,
-      OwnershipType ownershipType)
+      OwnershipType ownershipType,
+      Urn ownershipTypeUrn)
       throws Exception {
     final List<MetadataChangeProposal> changes =
-        buildAddOwnersProposals(opContext, ownerUrns, resources, ownershipType);
+        buildAddOwnersProposals(opContext, ownerUrns, resources, ownershipType, ownershipTypeUrn);
     ingestChangeProposals(opContext, changes);
   }
 
@@ -103,7 +104,8 @@ public class OwnerService extends BaseService {
       @Nonnull OperationContext opContext,
       List<com.linkedin.common.urn.Urn> ownerUrns,
       List<ResourceReference> resources,
-      OwnershipType ownershipType) {
+      OwnershipType ownershipType,
+      Urn ownershipTypeUrn) {
 
     final Map<Urn, Ownership> ownershipAspects =
         getOwnershipAspects(
@@ -127,7 +129,10 @@ public class OwnerService extends BaseService {
                 .setActor(
                     UrnUtils.getUrn(opContext.getSessionAuthentication().getActor().toUrnStr())));
       }
-      addOwnersIfNotExists(owners, ownerUrns, ownershipType);
+      ownerUrns.forEach(
+          ownerUrn -> {
+            OwnerServiceUtils.addOwnerToAspect(owners, ownerUrn, ownershipType, ownershipTypeUrn);
+          });
       proposals.add(
           buildMetadataChangeProposal(resource.getUrn(), Constants.OWNERSHIP_ASPECT_NAME, owners));
     }
@@ -160,43 +165,6 @@ public class OwnerService extends BaseService {
     return proposals;
   }
 
-  private void addOwnersIfNotExists(
-      Ownership owners, List<Urn> ownerUrns, OwnershipType ownershipType) {
-    if (!owners.hasOwners()) {
-      owners.setOwners(new OwnerArray());
-    }
-
-    OwnerArray ownerAssociationArray = owners.getOwners();
-
-    List<Urn> ownersToAdd = new ArrayList<>();
-    for (Urn ownerUrn : ownerUrns) {
-      if (ownerAssociationArray.stream()
-          .anyMatch(association -> association.getOwner().equals(ownerUrn))) {
-        continue;
-      }
-      ownersToAdd.add(ownerUrn);
-    }
-
-    // Check for no owners to add
-    if (ownersToAdd.size() == 0) {
-      return;
-    }
-
-    for (Urn ownerUrn : ownersToAdd) {
-      Owner newOwner = new Owner();
-      newOwner.setOwner(ownerUrn);
-      newOwner.setTypeUrn(mapOwnershipTypeToEntity(OwnershipType.NONE.name()));
-      newOwner.setType(ownershipType);
-      ownerAssociationArray.add(newOwner);
-    }
-  }
-
-  @VisibleForTesting
-  static Urn mapOwnershipTypeToEntity(String type) {
-    final String typeName = SYSTEM_ID + type.toLowerCase();
-    return Urn.createFromTuple(Constants.OWNERSHIP_TYPE_ENTITY_NAME, typeName);
-  }
-
   private static OwnerArray removeOwnersIfExists(Ownership owners, List<Urn> ownerUrns) {
     if (!owners.hasOwners()) {
       owners.setOwners(new OwnerArray());

metadata-service/services/src/main/java/com/linkedin/metadata/service/util/OwnerServiceUtils.java

@@ -0,0 +1,78 @@
+package com.linkedin.metadata.service.util;
+
+import com.linkedin.common.Owner;
+import com.linkedin.common.OwnerArray;
+import com.linkedin.common.Ownership;
+import com.linkedin.common.OwnershipSource;
+import com.linkedin.common.OwnershipSourceType;
+import com.linkedin.common.OwnershipType;
+import com.linkedin.common.urn.Urn;
+import com.linkedin.metadata.Constants;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+public class OwnerServiceUtils {
+  public static final String SYSTEM_ID = "__system__";
+
+  public static void addOwnerToAspect(
+      Ownership ownershipAspect, Urn ownerUrn, OwnershipType type, Urn ownershipTypeUrn) {
+    addOwnerToAspect(ownershipAspect, ownerUrn, type, ownershipTypeUrn, null);
+  }
+
+  public static void addOwnerToAspect(
+      Ownership ownershipAspect,
+      Urn ownerUrn,
+      OwnershipType type,
+      Urn ownershipTypeUrn,
+      @Nullable OwnershipSourceType sourceType) {
+    if (!ownershipAspect.hasOwners()) {
+      ownershipAspect.setOwners(new OwnerArray());
+    }
+
+    OwnerArray ownerArray = new OwnerArray(ownershipAspect.getOwners());
+    removeExistingOwnerIfExists(ownerArray, ownerUrn, ownershipTypeUrn);
+
+    Owner newOwner = new Owner();
+
+    newOwner.setType(type);
+    if (ownershipTypeUrn != null) {
+      newOwner.setTypeUrn(ownershipTypeUrn);
+    }
+    if (sourceType != null) {
+      newOwner.setSource(new OwnershipSource().setType(sourceType));
+    }
+    newOwner.setOwner(ownerUrn);
+    ownerArray.add(newOwner);
+    ownershipAspect.setOwners(ownerArray);
+  }
+
+  public static void removeExistingOwnerIfExists(
+      OwnerArray ownerArray, Urn ownerUrn, Urn ownershipTypeUrn) {
+    ownerArray.removeIf(
+        owner -> {
+          // Remove old ownership if it exists (check ownerUrn + type (entity & deprecated type))
+          return isOwnerEqual(owner, ownerUrn, ownershipTypeUrn);
+        });
+  }
+
+  public static boolean isOwnerEqual(
+      @Nonnull Owner owner, @Nonnull Urn ownerUrn, @Nullable Urn ownershipTypeUrn) {
+    if (!owner.getOwner().equals(ownerUrn)) {
+      return false;
+    }
+    if (owner.getTypeUrn() != null && ownershipTypeUrn != null) {
+      return owner.getTypeUrn().equals(ownershipTypeUrn);
+    }
+    if (ownershipTypeUrn == null) {
+      return true;
+    }
+    // Fall back to mapping deprecated type to the new ownership entity
+    return mapOwnershipTypeToEntity(OwnershipType.valueOf(owner.getType().toString()).name())
+        .equals(ownershipTypeUrn);
+  }
+
+  public static Urn mapOwnershipTypeToEntity(String type) {
+    final String typeName = SYSTEM_ID + type.toLowerCase();
+    return Urn.createFromTuple(Constants.OWNERSHIP_TYPE_ENTITY_NAME, typeName);
+  }
+}

metadata-service/services/src/test/java/com/linkedin/metadata/service/OwnerServiceTest.java

@@ -61,16 +61,14 @@ public class OwnerServiceTest {
             ImmutableList.of(
                 new ResourceReference(TEST_ENTITY_URN_1, null, null),
                 new ResourceReference(TEST_ENTITY_URN_2, null, null)),
-            OwnershipType.NONE);
+            OwnershipType.NONE,
+            null);
 
     OwnerArray expected =
         new OwnerArray(
             ImmutableList.of(
                 new Owner().setOwner(TEST_OWNER_URN_1).setType(OwnershipType.NONE),
-                new Owner()
+                new Owner().setOwner(newOwnerUrn).setType(OwnershipType.NONE)));
-                    .setOwner(newOwnerUrn)
-                    .setType(OwnershipType.NONE)
-                    .setTypeUrn(mapOwnershipTypeToEntity(OwnershipType.NONE.toString()))));
 
     MetadataChangeProposal event1 = events.get(0);
     Assert.assertEquals(event1.getAspectName(), Constants.OWNERSHIP_ASPECT_NAME);
@@ -103,15 +101,12 @@ public class OwnerServiceTest {
             ImmutableList.of(
                 new ResourceReference(TEST_ENTITY_URN_1, null, null),
                 new ResourceReference(TEST_ENTITY_URN_2, null, null)),
-            OwnershipType.NONE);
+            OwnershipType.NONE,
+            null);
 
     OwnerArray expectedOwners =
         new OwnerArray(
-            ImmutableList.of(
+            ImmutableList.of(new Owner().setOwner(newOwnerUrn).setType(OwnershipType.NONE)));
-                new Owner()
-                    .setOwner(newOwnerUrn)
-                    .setType(OwnershipType.NONE)
-                    .setTypeUrn(mapOwnershipTypeToEntity(OwnershipType.NONE.toString()))));
 
     MetadataChangeProposal event1 = events.get(0);
     Assert.assertEquals(event1.getAspectName(), Constants.OWNERSHIP_ASPECT_NAME);