Skip to content

Commit

Permalink
feat(forms) Handle deleting forms references when hard deleting forms (…
Browse files Browse the repository at this point in the history 
…datahub-project#10820)
  • Loading branch information
@chriscollins3456
chriscollins3456 authored Jul 1, 2024
1 parent 2d453c8 commit 64d4172
Show file tree
Hide file tree
Showing 3 changed files with 483 additions and 3 deletions.
214 changes: 213 additions & 1 deletion metadata-io/src/test/java/com/linkedin/metadata/entity/DeleteEntityServiceTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,10 @@
import com.datahub.util.RecordUtils;
import com.google.common.collect.ImmutableList;
import com.linkedin.common.AuditStamp;
import com.linkedin.common.FormAssociation;
import com.linkedin.common.FormAssociationArray;
import com.linkedin.common.FormVerificationAssociationArray;
import com.linkedin.common.Forms;
import com.linkedin.common.urn.Urn;
import com.linkedin.common.urn.UrnUtils;
import com.linkedin.container.Container;
Expand All @@ -19,11 +23,17 @@
import com.linkedin.metadata.event.EventProducer;
import com.linkedin.metadata.graph.GraphService;
import com.linkedin.metadata.graph.RelatedEntitiesResult;
import com.linkedin.metadata.query.filter.Filter;
import com.linkedin.metadata.query.filter.RelationshipDirection;
import com.linkedin.metadata.run.DeleteReferencesResponse;
import com.linkedin.metadata.search.EntitySearchService;
import com.linkedin.metadata.search.ScrollResult;
import com.linkedin.metadata.search.SearchEntity;
import com.linkedin.metadata.search.SearchEntityArray;
import com.linkedin.metadata.service.UpdateIndicesService;
import com.linkedin.metadata.utils.AuditStampUtils;
import com.linkedin.metadata.utils.SystemMetadataUtils;
import com.linkedin.mxe.MetadataChangeProposal;
import io.datahubproject.metadata.context.OperationContext;
import io.datahubproject.test.metadata.context.TestOperationContexts;
import java.sql.Timestamp;
Expand All @@ -39,17 +49,20 @@ public class DeleteEntityServiceTest {
protected GraphService _graphService = Mockito.mock(GraphService.class);
protected DeleteEntityService _deleteEntityService;
protected UpdateIndicesService _mockUpdateIndicesService;
protected EntitySearchService _mockSearchService;

public DeleteEntityServiceTest() {
opContext = TestOperationContexts.systemContextNoSearchAuthorization();
_aspectDao = mock(EbeanAspectDao.class);
_mockUpdateIndicesService = mock(UpdateIndicesService.class);
_mockSearchService = mock(EntitySearchService.class);
PreProcessHooks preProcessHooks = new PreProcessHooks();
preProcessHooks.setUiEnabled(true);
_entityServiceImpl =
new EntityServiceImpl(_aspectDao, mock(EventProducer.class), true, preProcessHooks, true);
_entityServiceImpl.setUpdateIndicesService(_mockUpdateIndicesService);
_deleteEntityService = new DeleteEntityService(_entityServiceImpl, _graphService);
_deleteEntityService =
new DeleteEntityService(_entityServiceImpl, _graphService, _mockSearchService);
}

/**
Expand Down Expand Up @@ -119,4 +132,203 @@ public void testDeleteUniqueRefGeneratesValidMCP() {
assertEquals(1, (int) response.getTotal());
assertFalse(response.getRelatedAspects().isEmpty());
}

/** This test checks whether updating search references works properly (for forms only for now) */
@Test
public void testDeleteSearchReferences() {
EntityService<?> mockEntityService = Mockito.mock(EntityService.class);
DeleteEntityService deleteEntityService =
new DeleteEntityService(mockEntityService, _graphService, _mockSearchService);

final Urn dataset = UrnUtils.toDatasetUrn("snowflake", "test", "DEV");
final Urn form = UrnUtils.getUrn("urn:li:form:12345");

ScrollResult scrollResult = new ScrollResult();
SearchEntityArray entities = new SearchEntityArray();
SearchEntity searchEntity = new SearchEntity();
searchEntity.setEntity(dataset);
entities.add(searchEntity);
scrollResult.setEntities(entities);
scrollResult.setNumEntities(1);
scrollResult.setScrollId("1");
Mockito.when(
_mockSearchService.structuredScroll(
Mockito.any(OperationContext.class),
Mockito.any(),
Mockito.eq("*"),
Mockito.any(Filter.class),
Mockito.eq(null),
Mockito.eq(null),
Mockito.eq("5m"),
Mockito.eq(1000)))
.thenReturn(scrollResult);

ScrollResult scrollResult2 = new ScrollResult();
scrollResult2.setNumEntities(0);
Mockito.when(
_mockSearchService.structuredScroll(
Mockito.any(OperationContext.class),
Mockito.any(),
Mockito.eq("*"),
Mockito.any(Filter.class),
Mockito.eq(null),
Mockito.eq("1"),
Mockito.eq("5m"),
Mockito.eq(1000)))
.thenReturn(scrollResult2);

Forms formsAspect = new Forms();
FormAssociationArray incompleteForms = new FormAssociationArray();
FormAssociation formAssociation = new FormAssociation();
formAssociation.setUrn(form);
incompleteForms.add(formAssociation);
formsAspect.setIncompleteForms(incompleteForms);
formsAspect.setCompletedForms(new FormAssociationArray());
formsAspect.setVerifications(new FormVerificationAssociationArray());
Mockito.when(
mockEntityService.getLatestAspect(
Mockito.any(OperationContext.class), Mockito.eq(dataset), Mockito.eq("forms")))
.thenReturn(formsAspect);

// no entities with relationships on forms
final RelatedEntitiesResult mockRelatedEntities =
new RelatedEntitiesResult(0, 0, 0, ImmutableList.of());
Mockito.when(
_graphService.findRelatedEntities(
null,
newFilter("urn", form.toString()),
null,
EMPTY_FILTER,
ImmutableList.of(),
newRelationshipFilter(EMPTY_FILTER, RelationshipDirection.INCOMING),
0,
10000))
.thenReturn(mockRelatedEntities);

final DeleteReferencesResponse response =
deleteEntityService.deleteReferencesTo(opContext, form, false);

// ensure we ingest one MCP for cleaning up forms reference
Mockito.verify(mockEntityService, Mockito.times(1))
.ingestProposal(
any(),
Mockito.any(MetadataChangeProposal.class),
Mockito.any(AuditStamp.class),
Mockito.eq(true));
assertEquals(1, (int) response.getTotal());
assertTrue(response.getRelatedAspects().isEmpty());
}

/** This test ensures we aren't issuing MCPs if there are no search references */
@Test
public void testDeleteNoSearchReferences() {
EntityService<?> mockEntityService = Mockito.mock(EntityService.class);
DeleteEntityService deleteEntityService =
new DeleteEntityService(mockEntityService, _graphService, _mockSearchService);

final Urn dataset = UrnUtils.toDatasetUrn("snowflake", "test", "DEV");
final Urn form = UrnUtils.getUrn("urn:li:form:12345");

ScrollResult scrollResult = new ScrollResult();
scrollResult.setEntities(new SearchEntityArray());
scrollResult.setNumEntities(0);
Mockito.when(
_mockSearchService.structuredScroll(
Mockito.any(OperationContext.class),
Mockito.any(),
Mockito.eq("*"),
Mockito.any(Filter.class),
Mockito.eq(null),
Mockito.eq(null),
Mockito.eq("5m"),
Mockito.eq(1000)))
.thenReturn(scrollResult);

// no entities with relationships on forms
final RelatedEntitiesResult mockRelatedEntities =
new RelatedEntitiesResult(0, 0, 0, ImmutableList.of());
Mockito.when(
_graphService.findRelatedEntities(
null,
newFilter("urn", form.toString()),
null,
EMPTY_FILTER,
ImmutableList.of(),
newRelationshipFilter(EMPTY_FILTER, RelationshipDirection.INCOMING),
0,
10000))
.thenReturn(mockRelatedEntities);

final DeleteReferencesResponse response =
deleteEntityService.deleteReferencesTo(opContext, form, false);

// ensure we did not ingest anything if there are no references
Mockito.verify(mockEntityService, Mockito.times(0))
.ingestProposal(
any(),
Mockito.any(MetadataChangeProposal.class),
Mockito.any(AuditStamp.class),
Mockito.eq(true));
assertEquals(0, (int) response.getTotal());
assertTrue(response.getRelatedAspects().isEmpty());
}

/** This test checks to make sure we don't issue MCPs if this is a dry-run */
@Test
public void testDeleteSearchReferencesDryRun() {
EntityService<?> mockEntityService = Mockito.mock(EntityService.class);
DeleteEntityService deleteEntityService =
new DeleteEntityService(mockEntityService, _graphService, _mockSearchService);

final Urn dataset = UrnUtils.toDatasetUrn("snowflake", "test", "DEV");
final Urn form = UrnUtils.getUrn("urn:li:form:12345");

ScrollResult scrollResult = new ScrollResult();
SearchEntityArray entities = new SearchEntityArray();
SearchEntity searchEntity = new SearchEntity();
searchEntity.setEntity(dataset);
entities.add(searchEntity);
scrollResult.setEntities(entities);
scrollResult.setNumEntities(1);
scrollResult.setScrollId("1");
Mockito.when(
_mockSearchService.structuredScroll(
Mockito.any(OperationContext.class),
Mockito.any(),
Mockito.eq("*"),
Mockito.any(Filter.class),
Mockito.eq(null),
Mockito.eq(null),
Mockito.eq("5m"),
Mockito.eq(1000)))
.thenReturn(scrollResult);

// no entities with relationships on forms
final RelatedEntitiesResult mockRelatedEntities =
new RelatedEntitiesResult(0, 0, 0, ImmutableList.of());
Mockito.when(
_graphService.findRelatedEntities(
null,
newFilter("urn", form.toString()),
null,
EMPTY_FILTER,
ImmutableList.of(),
newRelationshipFilter(EMPTY_FILTER, RelationshipDirection.INCOMING),
0,
10000))
.thenReturn(mockRelatedEntities);

final DeleteReferencesResponse response =
deleteEntityService.deleteReferencesTo(opContext, form, false);

// ensure we do not ingest anything since this is dry-run, but the total returns 1
Mockito.verify(mockEntityService, Mockito.times(0))
.ingestProposal(
any(),
Mockito.any(MetadataChangeProposal.class),
Mockito.any(AuditStamp.class),
Mockito.eq(true));
assertEquals(1, (int) response.getTotal());
assertTrue(response.getRelatedAspects().isEmpty());
}
}
7 changes: 6 additions & 1 deletion ...e/factories/src/main/java/com/linkedin/gms/factory/entity/DeleteEntityServiceFactory.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import com.linkedin.metadata.entity.DeleteEntityService;
import com.linkedin.metadata.entity.EntityService;
import com.linkedin.metadata.graph.GraphService;
import com.linkedin.metadata.search.EntitySearchService;
import javax.annotation.Nonnull;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
Expand All @@ -21,9 +22,13 @@ public class DeleteEntityServiceFactory {
@Qualifier("graphService")
private GraphService _graphService;

@Autowired
@Qualifier("entitySearchService")
private EntitySearchService _entitySearchService;

@Bean(name = "deleteEntityService")
@Nonnull
protected DeleteEntityService createDeleteEntityService() {
return new DeleteEntityService(_entityService, _graphService);
return new DeleteEntityService(_entityService, _graphService, _entitySearchService);
}
}
Loading

0 comments on commit 64d4172

Please sign in to comment.