diff --git a/metadata-io/src/test/java/com/linkedin/metadata/entity/DeleteEntityUtilsTest.java b/metadata-io/src/test/java/com/linkedin/metadata/entity/DeleteEntityUtilsTest.java index 943ad2967de429..f210ae94d0d70e 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/entity/DeleteEntityUtilsTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/entity/DeleteEntityUtilsTest.java @@ -1,14 +1,31 @@ package com.linkedin.metadata.entity; import com.datahub.util.RecordUtils; +import com.google.common.collect.ImmutableList; +import com.linkedin.common.urn.Urn; +import com.linkedin.common.urn.UrnUtils; import com.linkedin.data.DataList; import com.linkedin.data.DataMap; import com.linkedin.data.schema.DataSchema; import com.linkedin.data.schema.PathSpec; import com.linkedin.data.schema.grammar.PdlSchemaParser; import com.linkedin.data.schema.resolver.DefaultDataSchemaResolver; +import com.linkedin.data.template.StringArray; import com.linkedin.entity.Aspect; +import com.linkedin.form.FormInfo; +import com.linkedin.form.FormPrompt; +import com.linkedin.form.FormPromptArray; +import com.linkedin.form.StructuredPropertyParams; +import com.linkedin.metadata.query.filter.Condition; +import com.linkedin.metadata.query.filter.ConjunctiveCriterion; +import com.linkedin.metadata.query.filter.ConjunctiveCriterionArray; +import com.linkedin.metadata.query.filter.Criterion; +import com.linkedin.metadata.query.filter.CriterionArray; +import com.linkedin.metadata.query.filter.Filter; import com.linkedin.schema.SchemaMetadata; +import java.util.ArrayList; +import java.util.List; +import java.util.stream.Collectors; import junit.framework.TestCase; import org.testng.annotations.Test; @@ -359,4 +376,60 @@ public void testSchemaMetadataDelete() { .get("tags")) .size()); } + + @Test + public void testRemovePromptsFromFormInfo() { + Urn deletedPropertyUrn = UrnUtils.getUrn("urn:li:structuredProperty:1"); + Urn existingPropertyUrn = UrnUtils.getUrn("urn:li:structuredProperty:2"); + List prompts = new ArrayList<>(); + prompts.add( + new FormPrompt() + .setId("1") + .setStructuredPropertyParams( + new StructuredPropertyParams().setUrn(deletedPropertyUrn))); + prompts.add( + new FormPrompt() + .setId("2") + .setStructuredPropertyParams( + new StructuredPropertyParams().setUrn(existingPropertyUrn))); + FormInfo formInfo = new FormInfo().setPrompts(new FormPromptArray(prompts)); + + FormInfo updatedFormInfo = + DeleteEntityUtils.removePromptsFromFormInfoAspect(formInfo, deletedPropertyUrn); + + assertEquals(updatedFormInfo.getPrompts().size(), 1); + assertEquals( + updatedFormInfo.getPrompts(), + formInfo.getPrompts().stream() + .filter(prompt -> !prompt.getId().equals("1")) + .collect(Collectors.toList())); + } + + @Test + public void testFilterForStructuredPropDeletion() { + Urn deletedPropertyUrn = UrnUtils.getUrn("urn:li:structuredProperty:1"); + + final CriterionArray criterionArray = new CriterionArray(); + criterionArray.add( + new Criterion() + .setField("structuredPropertyPromptUrns") + .setValues(new StringArray(deletedPropertyUrn.toString())) + .setNegated(false) + .setValue("") + .setCondition(Condition.EQUAL)); + Filter expectedFilter = + new Filter() + .setOr( + new ConjunctiveCriterionArray(new ConjunctiveCriterion().setAnd(criterionArray))); + + assertEquals( + DeleteEntityUtils.getFilterForStructuredPropertyDeletion(deletedPropertyUrn), + expectedFilter); + } + + @Test + public void testEntityNamesForStructuredPropDeletion() { + assertEquals( + DeleteEntityUtils.getEntityNamesForStructuredPropertyDeletion(), ImmutableList.of("form")); + } } diff --git a/metadata-models/src/main/pegasus/com/linkedin/form/FormPrompt.pdl b/metadata-models/src/main/pegasus/com/linkedin/form/FormPrompt.pdl index f84a8a719f07a4..55e1a2a8fc1b72 100644 --- a/metadata-models/src/main/pegasus/com/linkedin/form/FormPrompt.pdl +++ b/metadata-models/src/main/pegasus/com/linkedin/form/FormPrompt.pdl @@ -43,6 +43,10 @@ record FormPrompt { /** * The structured property that is required on this entity */ + @Searchable = { + "fieldType": "URN", + "fieldName": "structuredPropertyPromptUrns", + } urn: Urn } diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/factories/BootstrapManagerFactory.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/factories/BootstrapManagerFactory.java index ffc739e905cd68..3f8d24b9b61a93 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/factories/BootstrapManagerFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/factories/BootstrapManagerFactory.java @@ -18,6 +18,7 @@ import com.linkedin.metadata.boot.steps.RemoveClientIdAspectStep; import com.linkedin.metadata.boot.steps.RestoreColumnLineageIndices; import com.linkedin.metadata.boot.steps.RestoreDbtSiblingsIndices; +import com.linkedin.metadata.boot.steps.RestoreFormInfoIndicesStep; import com.linkedin.metadata.boot.steps.RestoreGlossaryIndices; import com.linkedin.metadata.boot.steps.WaitForSystemUpdateStep; import com.linkedin.metadata.entity.AspectMigrationsDao; @@ -110,6 +111,8 @@ protected BootstrapManager createInstance( final WaitForSystemUpdateStep waitForSystemUpdateStep = new WaitForSystemUpdateStep(_dataHubUpgradeKafkaListener, _configurationProvider); final IngestEntityTypesStep ingestEntityTypesStep = new IngestEntityTypesStep(_entityService); + final RestoreFormInfoIndicesStep restoreFormInfoIndicesStep = + new RestoreFormInfoIndicesStep(_entityService); final List finalSteps = new ArrayList<>( @@ -124,7 +127,8 @@ protected BootstrapManager createInstance( restoreDbtSiblingsIndices, indexDataPlatformsStep, restoreColumnLineageIndices, - ingestEntityTypesStep)); + ingestEntityTypesStep, + restoreFormInfoIndicesStep)); return new BootstrapManager(finalSteps); } diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/RestoreFormInfoIndicesStep.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/RestoreFormInfoIndicesStep.java new file mode 100644 index 00000000000000..c5dbea5965a66a --- /dev/null +++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/RestoreFormInfoIndicesStep.java @@ -0,0 +1,129 @@ +package com.linkedin.metadata.boot.steps; + +import com.linkedin.common.AuditStamp; +import com.linkedin.common.urn.Urn; +import com.linkedin.data.template.RecordTemplate; +import com.linkedin.events.metadata.ChangeType; +import com.linkedin.form.FormInfo; +import com.linkedin.metadata.Constants; +import com.linkedin.metadata.boot.UpgradeStep; +import com.linkedin.metadata.entity.EntityService; +import com.linkedin.metadata.entity.ListResult; +import com.linkedin.metadata.models.AspectSpec; +import com.linkedin.metadata.query.ExtraInfo; +import io.datahubproject.metadata.context.OperationContext; +import java.util.LinkedList; +import java.util.List; +import java.util.Objects; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Future; +import javax.annotation.Nonnull; +import lombok.extern.slf4j.Slf4j; + +@Slf4j +public class RestoreFormInfoIndicesStep extends UpgradeStep { + private static final String VERSION = "1"; + private static final String UPGRADE_ID = "restore-form-info-indices"; + private static final Integer BATCH_SIZE = 1000; + + public RestoreFormInfoIndicesStep(@Nonnull final EntityService entityService) { + super(entityService, VERSION, UPGRADE_ID); + } + + @Override + public void upgrade(@Nonnull OperationContext systemOperationContext) throws Exception { + final AuditStamp auditStamp = + new AuditStamp() + .setActor(Urn.createFromString(Constants.SYSTEM_ACTOR)) + .setTime(System.currentTimeMillis()); + + final int totalFormCount = getAndRestoreFormInfoIndices(systemOperationContext, 0, auditStamp); + int formCount = BATCH_SIZE; + while (formCount < totalFormCount) { + getAndRestoreFormInfoIndices(systemOperationContext, formCount, auditStamp); + formCount += BATCH_SIZE; + } + } + + @Nonnull + @Override + public ExecutionMode getExecutionMode() { + return ExecutionMode.ASYNC; + } + + private int getAndRestoreFormInfoIndices( + @Nonnull OperationContext systemOperationContext, int start, AuditStamp auditStamp) { + final AspectSpec formInfoAspectSpec = + systemOperationContext + .getEntityRegistry() + .getEntitySpec(Constants.FORM_ENTITY_NAME) + .getAspectSpec(Constants.FORM_INFO_ASPECT_NAME); + + final ListResult latestAspects = + entityService.listLatestAspects( + systemOperationContext, + Constants.FORM_ENTITY_NAME, + Constants.FORM_INFO_ASPECT_NAME, + start, + BATCH_SIZE); + + if (latestAspects.getTotalCount() == 0 + || latestAspects.getValues() == null + || latestAspects.getMetadata() == null) { + log.debug("Found 0 formInfo aspects for forms. Skipping migration."); + return 0; + } + + if (latestAspects.getValues().size() != latestAspects.getMetadata().getExtraInfos().size()) { + // Bad result -- we should log that we cannot migrate this batch of formInfos. + log.warn( + "Failed to match formInfo aspects with corresponding urns. Found mismatched length between aspects ({})" + + "and metadata ({}) for metadata {}", + latestAspects.getValues().size(), + latestAspects.getMetadata().getExtraInfos().size(), + latestAspects.getMetadata()); + return latestAspects.getTotalCount(); + } + + List> futures = new LinkedList<>(); + for (int i = 0; i < latestAspects.getValues().size(); i++) { + ExtraInfo info = latestAspects.getMetadata().getExtraInfos().get(i); + RecordTemplate formInfoRecord = latestAspects.getValues().get(i); + Urn urn = info.getUrn(); + FormInfo formInfo = (FormInfo) formInfoRecord; + if (formInfo == null) { + log.warn("Received null formInfo for urn {}", urn); + continue; + } + + futures.add( + entityService + .alwaysProduceMCLAsync( + systemOperationContext, + urn, + Constants.FORM_ENTITY_NAME, + Constants.FORM_INFO_ASPECT_NAME, + formInfoAspectSpec, + null, + formInfo, + null, + null, + auditStamp, + ChangeType.RESTATE) + .getFirst()); + } + + futures.stream() + .filter(Objects::nonNull) + .forEach( + f -> { + try { + f.get(); + } catch (InterruptedException | ExecutionException e) { + throw new RuntimeException(e); + } + }); + + return latestAspects.getTotalCount(); + } +} diff --git a/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/RestoreFormInfoIndicesStepTest.java b/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/RestoreFormInfoIndicesStepTest.java new file mode 100644 index 00000000000000..9705ecbfcb739a --- /dev/null +++ b/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/RestoreFormInfoIndicesStepTest.java @@ -0,0 +1,242 @@ +package com.linkedin.metadata.boot.steps; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +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.entity.Aspect; +import com.linkedin.entity.EntityResponse; +import com.linkedin.entity.EnvelopedAspect; +import com.linkedin.entity.EnvelopedAspectMap; +import com.linkedin.events.metadata.ChangeType; +import com.linkedin.form.FormInfo; +import com.linkedin.metadata.Constants; +import com.linkedin.metadata.entity.EntityService; +import com.linkedin.metadata.entity.ListResult; +import com.linkedin.metadata.models.AspectSpec; +import com.linkedin.metadata.models.EntitySpec; +import com.linkedin.metadata.models.registry.EntityRegistry; +import com.linkedin.metadata.query.ExtraInfo; +import com.linkedin.metadata.query.ExtraInfoArray; +import com.linkedin.metadata.query.ListResultMetadata; +import com.linkedin.mxe.MetadataChangeProposal; +import com.linkedin.util.Pair; +import io.datahubproject.metadata.context.OperationContext; +import jakarta.annotation.Nonnull; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.Future; +import org.mockito.Mockito; +import org.testng.annotations.Test; + +public class RestoreFormInfoIndicesStepTest { + + private static final String VERSION_1 = "1"; + private static final String VERSION_2 = "2"; + private static final String FORM_INFO_UPGRADE_URN = + String.format( + "urn:li:%s:%s", Constants.DATA_HUB_UPGRADE_ENTITY_NAME, "restore-form-info-indices"); + private final Urn formUrn = UrnUtils.getUrn("urn:li:form:test"); + + @Test + public void testExecuteFirstTime() throws Exception { + final EntityService mockService = Mockito.mock(EntityService.class); + final EntityRegistry mockRegistry = Mockito.mock(EntityRegistry.class); + final OperationContext mockContext = mock(OperationContext.class); + when(mockContext.getEntityRegistry()).thenReturn(mockRegistry); + + mockGetUpgradeStep(mockContext, false, VERSION_1, mockService); + mockGetFormInfo(mockContext, formUrn, mockService); + + final AspectSpec aspectSpec = mockAspectSpecs(mockRegistry); + + final RestoreFormInfoIndicesStep restoreIndicesStep = + new RestoreFormInfoIndicesStep(mockService); + restoreIndicesStep.execute(mockContext); + + Mockito.verify(mockRegistry, Mockito.times(1)).getEntitySpec(Constants.FORM_ENTITY_NAME); + // creates upgradeRequest and upgradeResult aspects + Mockito.verify(mockService, Mockito.times(2)) + .ingestProposal( + any(OperationContext.class), + any(MetadataChangeProposal.class), + any(AuditStamp.class), + Mockito.eq(false)); + Mockito.verify(mockService, Mockito.times(1)) + .alwaysProduceMCLAsync( + any(OperationContext.class), + Mockito.eq(formUrn), + Mockito.eq(Constants.FORM_ENTITY_NAME), + Mockito.eq(Constants.FORM_INFO_ASPECT_NAME), + Mockito.eq(aspectSpec), + Mockito.eq(null), + any(), + Mockito.eq(null), + Mockito.eq(null), + any(), + Mockito.eq(ChangeType.RESTATE)); + } + + @Test + public void testExecuteWithNewVersion() throws Exception { + final EntityService mockService = Mockito.mock(EntityService.class); + final EntityRegistry mockRegistry = Mockito.mock(EntityRegistry.class); + final OperationContext mockContext = mock(OperationContext.class); + when(mockContext.getEntityRegistry()).thenReturn(mockRegistry); + + mockGetUpgradeStep(mockContext, true, VERSION_2, mockService); + mockGetFormInfo(mockContext, formUrn, mockService); + + final AspectSpec aspectSpec = mockAspectSpecs(mockRegistry); + + final RestoreFormInfoIndicesStep restoreIndicesStep = + new RestoreFormInfoIndicesStep(mockService); + restoreIndicesStep.execute(mockContext); + + Mockito.verify(mockRegistry, Mockito.times(1)).getEntitySpec(Constants.FORM_ENTITY_NAME); + // creates upgradeRequest and upgradeResult aspects + Mockito.verify(mockService, Mockito.times(2)) + .ingestProposal( + any(OperationContext.class), + any(MetadataChangeProposal.class), + any(AuditStamp.class), + Mockito.eq(false)); + Mockito.verify(mockService, Mockito.times(1)) + .alwaysProduceMCLAsync( + any(OperationContext.class), + Mockito.eq(formUrn), + Mockito.eq(Constants.FORM_ENTITY_NAME), + Mockito.eq(Constants.FORM_INFO_ASPECT_NAME), + Mockito.eq(aspectSpec), + Mockito.eq(null), + any(), + Mockito.eq(null), + Mockito.eq(null), + any(), + Mockito.eq(ChangeType.RESTATE)); + } + + @Test + public void testDoesNotExecuteWithSameVersion() throws Exception { + final EntityService mockService = Mockito.mock(EntityService.class); + final EntityRegistry mockRegistry = Mockito.mock(EntityRegistry.class); + final OperationContext mockContext = mock(OperationContext.class); + when(mockContext.getEntityRegistry()).thenReturn(mockRegistry); + + mockGetUpgradeStep(mockContext, true, VERSION_1, mockService); + mockGetFormInfo(mockContext, formUrn, mockService); + + final AspectSpec aspectSpec = mockAspectSpecs(mockRegistry); + + final RestoreFormInfoIndicesStep restoreIndicesStep = + new RestoreFormInfoIndicesStep(mockService); + restoreIndicesStep.execute(mockContext); + + Mockito.verify(mockRegistry, Mockito.times(0)).getEntitySpec(Constants.FORM_ENTITY_NAME); + // creates upgradeRequest and upgradeResult aspects + Mockito.verify(mockService, Mockito.times(0)) + .ingestProposal( + any(OperationContext.class), + any(MetadataChangeProposal.class), + any(AuditStamp.class), + Mockito.eq(false)); + Mockito.verify(mockService, Mockito.times(0)) + .alwaysProduceMCLAsync( + any(OperationContext.class), + Mockito.eq(formUrn), + Mockito.eq(Constants.FORM_ENTITY_NAME), + Mockito.eq(Constants.FORM_INFO_ASPECT_NAME), + Mockito.eq(aspectSpec), + Mockito.eq(null), + any(), + Mockito.eq(null), + Mockito.eq(null), + any(), + Mockito.eq(ChangeType.RESTATE)); + } + + private void mockGetFormInfo( + @Nonnull OperationContext mockContext, + @Nonnull Urn formUrn, + @Nonnull EntityService mockService) { + final List extraInfos = + ImmutableList.of( + new ExtraInfo() + .setUrn(formUrn) + .setVersion(0L) + .setAudit( + new AuditStamp() + .setActor(UrnUtils.getUrn("urn:li:corpuser:test")) + .setTime(0L))); + + when(mockService.alwaysProduceMCLAsync( + any(OperationContext.class), + any(Urn.class), + Mockito.anyString(), + Mockito.anyString(), + any(AspectSpec.class), + Mockito.eq(null), + any(), + any(), + any(), + any(), + any(ChangeType.class))) + .thenReturn(Pair.of(Mockito.mock(Future.class), false)); + + when(mockService.listLatestAspects( + any(OperationContext.class), + Mockito.eq(Constants.FORM_ENTITY_NAME), + Mockito.eq(Constants.FORM_INFO_ASPECT_NAME), + Mockito.eq(0), + Mockito.eq(1000))) + .thenReturn( + new ListResult<>( + ImmutableList.of(new FormInfo()), + new ListResultMetadata().setExtraInfos(new ExtraInfoArray(extraInfos)), + 1, + false, + 1, + 1, + 1)); + } + + private AspectSpec mockAspectSpecs(@Nonnull EntityRegistry mockRegistry) { + final EntitySpec entitySpec = Mockito.mock(EntitySpec.class); + final AspectSpec aspectSpec = Mockito.mock(AspectSpec.class); + // Mock for formInfo + when(mockRegistry.getEntitySpec(Constants.FORM_ENTITY_NAME)).thenReturn(entitySpec); + when(entitySpec.getAspectSpec(Constants.FORM_INFO_ASPECT_NAME)).thenReturn(aspectSpec); + + return aspectSpec; + } + + private void mockGetUpgradeStep( + @Nonnull OperationContext mockContext, + boolean shouldReturnResponse, + @Nonnull String version, + @Nonnull EntityService mockService) + throws Exception { + + final Urn upgradeEntityUrn = UrnUtils.getUrn(FORM_INFO_UPGRADE_URN); + final com.linkedin.upgrade.DataHubUpgradeRequest upgradeRequest = + new com.linkedin.upgrade.DataHubUpgradeRequest().setVersion(version); + final Map upgradeRequestAspects = new HashMap<>(); + upgradeRequestAspects.put( + Constants.DATA_HUB_UPGRADE_REQUEST_ASPECT_NAME, + new EnvelopedAspect().setValue(new Aspect(upgradeRequest.data()))); + final EntityResponse response = + new EntityResponse().setAspects(new EnvelopedAspectMap(upgradeRequestAspects)); + when(mockService.getEntityV2( + mockContext, + Constants.DATA_HUB_UPGRADE_ENTITY_NAME, + upgradeEntityUrn, + Collections.singleton(Constants.DATA_HUB_UPGRADE_REQUEST_ASPECT_NAME))) + .thenReturn(shouldReturnResponse ? response : null); + } +} diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/entity/DeleteEntityService.java b/metadata-service/services/src/main/java/com/linkedin/metadata/entity/DeleteEntityService.java index a3c57a19eddd55..2259b9365ff640 100644 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/entity/DeleteEntityService.java +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/entity/DeleteEntityService.java @@ -1,16 +1,11 @@ package com.linkedin.metadata.entity; import static com.linkedin.metadata.search.utils.QueryUtils.*; -import static com.linkedin.metadata.utils.CriterionUtils.buildCriterion; import com.datahub.util.RecordUtils; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.linkedin.common.AuditStamp; -import com.linkedin.common.FormAssociation; -import com.linkedin.common.FormAssociationArray; -import com.linkedin.common.FormVerificationAssociation; -import com.linkedin.common.FormVerificationAssociationArray; import com.linkedin.common.Forms; import com.linkedin.common.urn.Urn; import com.linkedin.common.urn.UrnUtils; @@ -20,6 +15,7 @@ import com.linkedin.entity.EntityResponse; import com.linkedin.entity.EnvelopedAspect; import com.linkedin.events.metadata.ChangeType; +import com.linkedin.form.FormInfo; import com.linkedin.metadata.Constants; import com.linkedin.metadata.aspect.models.graph.RelatedEntity; import com.linkedin.metadata.graph.GraphService; @@ -28,10 +24,6 @@ import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.models.RelationshipFieldSpec; import com.linkedin.metadata.models.extractor.FieldExtractor; -import com.linkedin.metadata.query.filter.Condition; -import com.linkedin.metadata.query.filter.ConjunctiveCriterion; -import com.linkedin.metadata.query.filter.ConjunctiveCriterionArray; -import com.linkedin.metadata.query.filter.CriterionArray; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.query.filter.RelationshipDirection; import com.linkedin.metadata.run.DeleteReferencesResponse; @@ -568,65 +560,44 @@ private AssetScrollResult getAssetsReferencingUrn( result.assets = new ArrayList<>(); if (deletedUrn.getEntityType().equals("form")) { - // first, get all entities with this form assigned on it - final CriterionArray incompleteFormsArray = new CriterionArray(); - incompleteFormsArray.add( - buildCriterion("incompleteForms", Condition.EQUAL, deletedUrn.toString())); - final CriterionArray completedFormsArray = new CriterionArray(); - completedFormsArray.add( - buildCriterion("completedForms", Condition.EQUAL, deletedUrn.toString())); - // next, get all metadata tests created for this form - final CriterionArray metadataTestSourceArray = new CriterionArray(); - metadataTestSourceArray.add( - buildCriterion("sourceEntity", Condition.EQUAL, deletedUrn.toString())); - metadataTestSourceArray.add(buildCriterion("sourceType", Condition.EQUAL, "FORMS")); - Filter filter = - new Filter() - .setOr( - new ConjunctiveCriterionArray( - new ConjunctiveCriterion().setAnd(incompleteFormsArray), - new ConjunctiveCriterion().setAnd(completedFormsArray), - new ConjunctiveCriterion().setAnd(metadataTestSourceArray))); - ScrollResult scrollResult = - _searchService.structuredScroll( - opContext, - ImmutableList.of( - "dataset", - "dataJob", - "dataFlow", - "chart", - "dashboard", - "corpuser", - "corpGroup", - "domain", - "container", - "glossaryTerm", - "glossaryNode", - "mlModel", - "mlModelGroup", - "mlFeatureTable", - "mlFeature", - "mlPrimaryKey", - "schemaField", - "dataProduct", - "test"), - "*", - filter, - null, - scrollId, - "5m", - dryRun ? 1 : BATCH_SIZE); // need to pass in 1 for count otherwise get index error - if (scrollResult.getNumEntities() == 0 || scrollResult.getEntities().size() == 0) { - return result; - } - result.scrollId = scrollResult.getScrollId(); - result.totalAssetCount = scrollResult.getNumEntities(); - result.assets = - scrollResult.getEntities().stream() - .map(SearchEntity::getEntity) - .collect(Collectors.toList()); + Filter filter = DeleteEntityUtils.getFilterForFormDeletion(deletedUrn); + List entityNames = DeleteEntityUtils.getEntityNamesForFormDeletion(); + return scrollForAssets(opContext, result, filter, entityNames, scrollId, dryRun); + } + if (deletedUrn.getEntityType().equals("structuredProperty")) { + Filter filter = DeleteEntityUtils.getFilterForStructuredPropertyDeletion(deletedUrn); + List entityNames = DeleteEntityUtils.getEntityNamesForStructuredPropertyDeletion(); + return scrollForAssets(opContext, result, filter, entityNames, scrollId, dryRun); + } + return result; + } + + private AssetScrollResult scrollForAssets( + @Nonnull OperationContext opContext, + AssetScrollResult result, + final @Nullable Filter filter, + final @Nonnull List entityNames, + @Nullable String scrollId, + final boolean dryRun) { + ScrollResult scrollResult = + _searchService.structuredScroll( + opContext, + entityNames, + "*", + filter, + null, + scrollId, + "5m", + dryRun ? 1 : BATCH_SIZE); // need to pass in 1 for count otherwise get index error + if (scrollResult.getNumEntities() == 0 || scrollResult.getEntities().size() == 0) { return result; } + result.scrollId = scrollResult.getScrollId(); + result.totalAssetCount = scrollResult.getNumEntities(); + result.assets = + scrollResult.getEntities().stream() + .map(SearchEntity::getEntity) + .collect(Collectors.toList()); return result; } @@ -641,7 +612,7 @@ private List deleteSearchReferencesForAsset( } List mcps = new ArrayList<>(); - List aspectsToUpdate = getAspectsToUpdate(deletedUrn); + List aspectsToUpdate = getAspectsToUpdate(deletedUrn, assetUrn); aspectsToUpdate.forEach( aspectName -> { try { @@ -667,10 +638,15 @@ private List deleteSearchReferencesForAsset( *

TODO: extend this to support other types of deletes and be more dynamic depending on aspects * that the asset has */ - private List getAspectsToUpdate(@Nonnull final Urn deletedUrn) { + private List getAspectsToUpdate( + @Nonnull final Urn deletedUrn, @Nonnull final Urn assetUrn) { if (deletedUrn.getEntityType().equals("form")) { return ImmutableList.of("forms"); } + if (deletedUrn.getEntityType().equals("structuredProperty") + && assetUrn.getEntityType().equals("form")) { + return ImmutableList.of("formInfo"); + } return new ArrayList<>(); } @@ -697,6 +673,9 @@ private MetadataChangeProposal updateAspectForSearchReference( if (aspectName.equals("forms")) { return updateFormsAspect(opContext, assetUrn, deletedUrn); } + if (aspectName.equals("formInfo") && deletedUrn.getEntityType().equals("structuredProperty")) { + return updateFormInfoAspect(opContext, assetUrn, deletedUrn); + } return null; } @@ -708,35 +687,20 @@ private MetadataChangeProposal updateFormsAspect( RecordTemplate record = _entityService.getLatestAspect(opContext, assetUrn, "forms"); if (record == null) return null; - Forms formsAspect = new Forms(record.data()); - final AtomicReference updatedAspect; - try { - updatedAspect = new AtomicReference<>(formsAspect.copy()); - } catch (Exception e) { - throw new RuntimeException("Failed to copy the forms aspect for updating", e); - } - - List incompleteForms = - formsAspect.getIncompleteForms().stream() - .filter(incompleteForm -> !incompleteForm.getUrn().equals(deletedUrn)) - .collect(Collectors.toList()); - List completedForms = - formsAspect.getCompletedForms().stream() - .filter(completedForm -> !completedForm.getUrn().equals(deletedUrn)) - .collect(Collectors.toList()); - final List verifications = - formsAspect.getVerifications().stream() - .filter(verification -> !verification.getForm().equals(deletedUrn)) - .collect(Collectors.toList()); + return DeleteEntityUtils.removeFormFromFormsAspect( + new Forms(record.data()), assetUrn, deletedUrn); + } - updatedAspect.get().setIncompleteForms(new FormAssociationArray(incompleteForms)); - updatedAspect.get().setCompletedForms(new FormAssociationArray(completedForms)); - updatedAspect.get().setVerifications(new FormVerificationAssociationArray(verifications)); + @Nullable + private MetadataChangeProposal updateFormInfoAspect( + @Nonnull OperationContext opContext, + @Nonnull final Urn assetUrn, + @Nonnull final Urn deletedUrn) { + RecordTemplate record = _entityService.getLatestAspect(opContext, assetUrn, "formInfo"); + if (record == null) return null; - if (!formsAspect.equals(updatedAspect.get())) { - return AspectUtils.buildMetadataChangeProposal(assetUrn, "forms", updatedAspect.get()); - } - return null; + return DeleteEntityUtils.createFormInfoUpdateProposal( + new FormInfo(record.data()), assetUrn, deletedUrn); } private AuditStamp createAuditStamp() { diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/entity/DeleteEntityUtils.java b/metadata-service/services/src/main/java/com/linkedin/metadata/entity/DeleteEntityUtils.java index 0a8b5880e5bce9..20dc104e1b436e 100644 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/entity/DeleteEntityUtils.java +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/entity/DeleteEntityUtils.java @@ -1,5 +1,14 @@ package com.linkedin.metadata.entity; +import static com.linkedin.metadata.utils.CriterionUtils.buildCriterion; + +import com.google.common.collect.ImmutableList; +import com.linkedin.common.FormAssociation; +import com.linkedin.common.FormAssociationArray; +import com.linkedin.common.FormVerificationAssociation; +import com.linkedin.common.FormVerificationAssociationArray; +import com.linkedin.common.Forms; +import com.linkedin.common.urn.Urn; import com.linkedin.data.DataComplex; import com.linkedin.data.DataList; import com.linkedin.data.DataMap; @@ -9,8 +18,22 @@ import com.linkedin.data.schema.RecordDataSchema; import com.linkedin.data.template.RecordTemplate; import com.linkedin.entity.Aspect; +import com.linkedin.form.FormInfo; +import com.linkedin.form.FormPrompt; +import com.linkedin.form.FormPromptArray; +import com.linkedin.metadata.query.filter.Condition; +import com.linkedin.metadata.query.filter.ConjunctiveCriterion; +import com.linkedin.metadata.query.filter.ConjunctiveCriterionArray; +import com.linkedin.metadata.query.filter.CriterionArray; +import com.linkedin.metadata.query.filter.Filter; +import com.linkedin.metadata.utils.CriterionUtils; +import com.linkedin.mxe.MetadataChangeProposal; import java.util.List; import java.util.ListIterator; +import java.util.concurrent.atomic.AtomicReference; +import java.util.stream.Collectors; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; import lombok.extern.slf4j.Slf4j; /** @@ -207,4 +230,147 @@ private static DataComplex removeValueFromArray( } return aspectList; } + + /* + * Form Deletion Section + */ + + // We need to update assets that have this form on them in one way or another + public static Filter getFilterForFormDeletion(@Nonnull final Urn deletedUrn) { + // first, get all entities with this form assigned on it + final CriterionArray incompleteFormsArray = new CriterionArray(); + incompleteFormsArray.add( + buildCriterion("incompleteForms", Condition.EQUAL, deletedUrn.toString())); + final CriterionArray completedFormsArray = new CriterionArray(); + completedFormsArray.add( + buildCriterion("completedForms", Condition.EQUAL, deletedUrn.toString())); + // next, get all metadata tests created for this form + final CriterionArray metadataTestSourceArray = new CriterionArray(); + metadataTestSourceArray.add( + buildCriterion("sourceEntity", Condition.EQUAL, deletedUrn.toString())); + metadataTestSourceArray.add(buildCriterion("sourceType", Condition.EQUAL, "FORMS")); + return new Filter() + .setOr( + new ConjunctiveCriterionArray( + new ConjunctiveCriterion().setAnd(incompleteFormsArray), + new ConjunctiveCriterion().setAnd(completedFormsArray), + new ConjunctiveCriterion().setAnd(metadataTestSourceArray))); + } + + @Nullable + public static MetadataChangeProposal removeFormFromFormsAspect( + @Nonnull Forms formsAspect, @Nonnull final Urn assetUrn, @Nonnull final Urn deletedUrn) { + final AtomicReference updatedAspect; + try { + updatedAspect = new AtomicReference<>(formsAspect.copy()); + } catch (Exception e) { + throw new RuntimeException("Failed to copy the forms aspect for updating", e); + } + + List incompleteForms = + formsAspect.getIncompleteForms().stream() + .filter(incompleteForm -> !incompleteForm.getUrn().equals(deletedUrn)) + .collect(Collectors.toList()); + List completedForms = + formsAspect.getCompletedForms().stream() + .filter(completedForm -> !completedForm.getUrn().equals(deletedUrn)) + .collect(Collectors.toList()); + final List verifications = + formsAspect.getVerifications().stream() + .filter(verification -> !verification.getForm().equals(deletedUrn)) + .collect(Collectors.toList()); + + updatedAspect.get().setIncompleteForms(new FormAssociationArray(incompleteForms)); + updatedAspect.get().setCompletedForms(new FormAssociationArray(completedForms)); + updatedAspect.get().setVerifications(new FormVerificationAssociationArray(verifications)); + + if (!formsAspect.equals(updatedAspect.get())) { + return AspectUtils.buildMetadataChangeProposal(assetUrn, "forms", updatedAspect.get()); + } + return null; + } + + // all assets that could have a form associated with them + public static List getEntityNamesForFormDeletion() { + return ImmutableList.of( + "dataset", + "dataJob", + "dataFlow", + "chart", + "dashboard", + "corpuser", + "corpGroup", + "domain", + "container", + "glossaryTerm", + "glossaryNode", + "mlModel", + "mlModelGroup", + "mlFeatureTable", + "mlFeature", + "mlPrimaryKey", + "schemaField", + "dataProduct", + "test"); + } + + /* + * Structured Property Deletion Section + */ + + // get forms that have this structured property referenced in a prompt + public static Filter getFilterForStructuredPropertyDeletion(@Nonnull final Urn deletedUrn) { + final CriterionArray criterionArray = new CriterionArray(); + criterionArray.add( + CriterionUtils.buildCriterion( + "structuredPropertyPromptUrns", Condition.EQUAL, deletedUrn.toString())); + return new Filter() + .setOr(new ConjunctiveCriterionArray(new ConjunctiveCriterion().setAnd(criterionArray))); + } + + // only need to update forms manually when deleting structured props + public static List getEntityNamesForStructuredPropertyDeletion() { + return ImmutableList.of("form"); + } + + @Nullable + public static MetadataChangeProposal createFormInfoUpdateProposal( + @Nonnull FormInfo formsAspect, @Nonnull final Urn assetUrn, @Nonnull final Urn deletedUrn) { + final FormInfo updatedFormInfo = removePromptsFromFormInfoAspect(formsAspect, deletedUrn); + + if (!formsAspect.equals(updatedFormInfo)) { + return AspectUtils.buildMetadataChangeProposal(assetUrn, "formInfo", updatedFormInfo); + } + + return null; + } + + // remove any prompts referencing the deleted structured property urn + @Nonnull + public static FormInfo removePromptsFromFormInfoAspect( + @Nonnull FormInfo formsAspect, @Nonnull final Urn deletedUrn) { + final AtomicReference updatedAspect; + try { + updatedAspect = new AtomicReference<>(formsAspect.copy()); + } catch (Exception e) { + throw new RuntimeException("Failed to copy the formInfo aspect for updating", e); + } + + // filter out any prompt that has this structured property referenced on it + List filteredPrompts = + formsAspect.getPrompts().stream() + .filter( + prompt -> { + if (prompt.getStructuredPropertyParams() != null + && prompt.getStructuredPropertyParams().getUrn() != null) { + return !prompt.getStructuredPropertyParams().getUrn().equals(deletedUrn); + } + return true; + }) + .collect(Collectors.toList()); + + updatedAspect.get().setPrompts(new FormPromptArray(filteredPrompts)); + + return updatedAspect.get(); + } }