fix(gms): Fixes delete references for single relationship aspects

metadata-io/src/main/java/com/linkedin/metadata/entity/DeleteEntityService.java

@@ -28,6 +28,7 @@
 import com.linkedin.mxe.MetadataChangeProposal;
 
 import java.net.URISyntaxException;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -241,21 +242,14 @@ private void deleteReference(final Urn urn, final RelatedEntity relatedEntity) {
      * @param prevAspect the old value for the aspect
      */
     private void deleteAspect(Urn urn, String aspectName, RecordTemplate prevAspect) {
-        final MetadataChangeProposal proposal = new MetadataChangeProposal();
-        proposal.setEntityUrn(urn);
-        proposal.setChangeType(ChangeType.DELETE);
-        proposal.setEntityType(urn.getEntityType());
-        proposal.setAspectName(aspectName);
-
-        final AuditStamp auditStamp = new AuditStamp().setActor(UrnUtils.getUrn(Constants.SYSTEM_ACTOR)).setTime(System.currentTimeMillis());
-        final EntityService.IngestProposalResult ingestProposalResult = _entityService.ingestProposal(proposal, auditStamp, false);
-
-        if (!ingestProposalResult.isDidUpdate()) {
-            log.error("Failed to ingest aspect with references removed. Before {}, after: null, please check MCP processor"
+        final RollbackResult rollbackResult = _entityService.deleteAspect(urn.toString(), aspectName,
+            new HashMap<>(), true);
+        if (rollbackResult == null || rollbackResult.getNewValue() != null) {
+            log.error("Failed to delete aspect with references. Before {}, after: null, please check GMS logs"
                 + " logs for more information", prevAspect);
             handleError(new DeleteEntityServiceError("Failed to ingest new aspect",
-                DeleteEntityServiceErrorReason.MCP_PROCESSOR_FAILED,
-                ImmutableMap.of("proposal", proposal)));
+                DeleteEntityServiceErrorReason.ASPECT_DELETE_FAILED,
+                ImmutableMap.of("urn", urn, "aspectName", aspectName)));
         }
     }
 
@@ -380,7 +374,6 @@ private Stream<RelationshipFieldSpec> findRelationshipFor(final AspectSpec spec,
         return spec.getRelationshipFieldSpecs().stream()
                 .filter(relationship -> relationship.getRelationshipName().equals(relationshipType)
                         && relationship.getValidDestinationTypes().contains(entityType));
-        //.collect(Collectors.toList());
     }
 
     /**
@@ -403,6 +396,7 @@ private enum DeleteEntityServiceErrorReason {
         ENTITY_SERVICE_ASPECT_NOT_FOUND,
         ENTITY_REGISTRY_SPEC_NOT_FOUND,
         MCP_PROCESSOR_FAILED,
+        ASPECT_DELETE_FAILED,
         CLONE_FAILED,
     }
 

metadata-io/src/test/java/com/linkedin/metadata/entity/DeleteEntityServiceTest.java

@@ -0,0 +1,100 @@
+package com.linkedin.metadata.entity;
+
+import com.datahub.util.RecordUtils;
+import com.google.common.collect.ImmutableList;
+import com.linkedin.common.AuditStamp;
+import com.linkedin.common.urn.Urn;
+import com.linkedin.common.urn.UrnUtils;
+import com.linkedin.container.Container;
+import com.linkedin.entity.EntityResponse;
+import com.linkedin.events.metadata.ChangeType;
+import com.linkedin.metadata.Constants;
+import com.linkedin.metadata.entity.ebean.EbeanAspectDao;
+import com.linkedin.metadata.event.EventProducer;
+import com.linkedin.metadata.graph.GraphService;
+import com.linkedin.metadata.graph.RelatedEntitiesResult;
+import com.linkedin.metadata.graph.RelatedEntity;
+import com.linkedin.metadata.models.registry.ConfigEntityRegistry;
+import com.linkedin.metadata.models.registry.EntityRegistry;
+import com.linkedin.metadata.query.filter.RelationshipDirection;
+import com.linkedin.metadata.run.DeleteReferencesResponse;
+import com.linkedin.metadata.snapshot.Snapshot;
+import com.linkedin.metadata.utils.AuditStampUtils;
+import com.linkedin.metadata.utils.SystemMetadataUtils;
+import java.sql.Timestamp;
+import java.util.Map;
+import org.mockito.Mockito;
+import org.testng.annotations.Test;
+
+import static com.linkedin.metadata.search.utils.QueryUtils.*;
+import static org.mockito.Mockito.*;
+import static org.testng.AssertJUnit.*;
+
+
+public class DeleteEntityServiceTest {
+
+  protected EbeanAspectDao _aspectDao;
+
+  protected EntityService _entityService;
+  protected GraphService _graphService = Mockito.mock(GraphService.class);
+
+  protected DeleteEntityService _deleteEntityService;
+
+  protected EntityRegistry _entityRegistry;
+
+  public DeleteEntityServiceTest() {
+    _entityRegistry = new ConfigEntityRegistry(Snapshot.class.getClassLoader()
+        .getResourceAsStream("entity-registry.yml"));
+    _aspectDao = mock(EbeanAspectDao.class);
+    _entityService = new EntityService(_aspectDao, mock(EventProducer.class), _entityRegistry);
+    _deleteEntityService = new DeleteEntityService(_entityService, _graphService);
+  }
+
+  /**
+   * This test checks whether deleting non array references in PDL aspects generates a valid MCP.
+   */
+  @Test
+  public void testDeleteUniqueRefGeneratesValidMCP() {
+    final Urn dataset = UrnUtils.toDatasetUrn("snowflake", "test", "DEV");
+    final Urn container = UrnUtils.getUrn("urn:li:container:d1006cf3-3ff9-48e3-85cd-26eb23775ab2");
+
+    final RelatedEntitiesResult mockRelatedEntities =
+        new RelatedEntitiesResult(0, 1, 1, ImmutableList.of(new RelatedEntity("IsPartOf", dataset.toString())));
+
+    Mockito.when(_graphService.findRelatedEntities(null, newFilter("urn", container.toString()),
+            null, EMPTY_FILTER, ImmutableList.of(),
+            newRelationshipFilter(EMPTY_FILTER, RelationshipDirection.INCOMING), 0, 10000))
+        .thenReturn(mockRelatedEntities);
+
+    final EntityResponse entityResponse = new EntityResponse();
+    entityResponse.setUrn(dataset);
+    entityResponse.setEntityName(dataset.getEntityType());
+    final Container containerAspect = new Container();
+    containerAspect.setContainer(container);
+    final EntityAspectIdentifier dbKey = new EntityAspectIdentifier(dataset.toString(), Constants.CONTAINER_ASPECT_NAME, 0);
+
+    final EntityAspect dbValue = new EntityAspect();
+    dbValue.setUrn(dataset.toString());
+    dbValue.setVersion(0);
+    dbValue.setAspect(Constants.CONTAINER_ASPECT_NAME);
+    dbValue.setMetadata(RecordUtils.toJsonString(containerAspect));
+    dbValue.setSystemMetadata(RecordUtils.toJsonString(SystemMetadataUtils.createDefaultSystemMetadata()));
+    final AuditStamp auditStamp = AuditStampUtils.createDefaultAuditStamp();
+    dbValue.setCreatedBy(auditStamp.getActor().toString());
+    dbValue.setCreatedOn(new Timestamp(auditStamp.getTime()));
+
+    final Map<EntityAspectIdentifier, EntityAspect> dbEntries = Map.of(dbKey, dbValue);
+    Mockito.when(_aspectDao.batchGet(Mockito.any())).thenReturn(dbEntries);
+
+    RollbackResult result = new RollbackResult(container, Constants.DATASET_ENTITY_NAME,
+        Constants.CONTAINER_ASPECT_NAME, containerAspect, null, null, null,
+        ChangeType.DELETE, false, 1);
+
+    Mockito.when(_aspectDao.runInTransactionWithRetry(Mockito.any(), Mockito.anyInt()))
+        .thenReturn(result);
+
+    final DeleteReferencesResponse response = _deleteEntityService.deleteReferencesTo(container, false);
+    assertEquals(1, (int) response.getTotal());
+    assertFalse(response.getRelatedAspects().isEmpty());
+  }
+}