Skip to content

Commit 96ce338

Browse files
chriscollins3456sleeperdeep
authored andcommitted
feat(forms) Clean up form prompts on structured property deletion (datahub-project#12053)
1 parent db9c2ee commit 96ce338

File tree

7 files changed

+678
-96
lines changed

7 files changed

+678
-96
lines changed

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

+73
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,31 @@
11
package com.linkedin.metadata.entity;
22

33
import com.datahub.util.RecordUtils;
4+
import com.google.common.collect.ImmutableList;
5+
import com.linkedin.common.urn.Urn;
6+
import com.linkedin.common.urn.UrnUtils;
47
import com.linkedin.data.DataList;
58
import com.linkedin.data.DataMap;
69
import com.linkedin.data.schema.DataSchema;
710
import com.linkedin.data.schema.PathSpec;
811
import com.linkedin.data.schema.grammar.PdlSchemaParser;
912
import com.linkedin.data.schema.resolver.DefaultDataSchemaResolver;
13+
import com.linkedin.data.template.StringArray;
1014
import com.linkedin.entity.Aspect;
15+
import com.linkedin.form.FormInfo;
16+
import com.linkedin.form.FormPrompt;
17+
import com.linkedin.form.FormPromptArray;
18+
import com.linkedin.form.StructuredPropertyParams;
19+
import com.linkedin.metadata.query.filter.Condition;
20+
import com.linkedin.metadata.query.filter.ConjunctiveCriterion;
21+
import com.linkedin.metadata.query.filter.ConjunctiveCriterionArray;
22+
import com.linkedin.metadata.query.filter.Criterion;
23+
import com.linkedin.metadata.query.filter.CriterionArray;
24+
import com.linkedin.metadata.query.filter.Filter;
1125
import com.linkedin.schema.SchemaMetadata;
26+
import java.util.ArrayList;
27+
import java.util.List;
28+
import java.util.stream.Collectors;
1229
import junit.framework.TestCase;
1330
import org.testng.annotations.Test;
1431

@@ -359,4 +376,60 @@ public void testSchemaMetadataDelete() {
359376
.get("tags"))
360377
.size());
361378
}
379+
380+
@Test
381+
public void testRemovePromptsFromFormInfo() {
382+
Urn deletedPropertyUrn = UrnUtils.getUrn("urn:li:structuredProperty:1");
383+
Urn existingPropertyUrn = UrnUtils.getUrn("urn:li:structuredProperty:2");
384+
List<FormPrompt> prompts = new ArrayList<>();
385+
prompts.add(
386+
new FormPrompt()
387+
.setId("1")
388+
.setStructuredPropertyParams(
389+
new StructuredPropertyParams().setUrn(deletedPropertyUrn)));
390+
prompts.add(
391+
new FormPrompt()
392+
.setId("2")
393+
.setStructuredPropertyParams(
394+
new StructuredPropertyParams().setUrn(existingPropertyUrn)));
395+
FormInfo formInfo = new FormInfo().setPrompts(new FormPromptArray(prompts));
396+
397+
FormInfo updatedFormInfo =
398+
DeleteEntityUtils.removePromptsFromFormInfoAspect(formInfo, deletedPropertyUrn);
399+
400+
assertEquals(updatedFormInfo.getPrompts().size(), 1);
401+
assertEquals(
402+
updatedFormInfo.getPrompts(),
403+
formInfo.getPrompts().stream()
404+
.filter(prompt -> !prompt.getId().equals("1"))
405+
.collect(Collectors.toList()));
406+
}
407+
408+
@Test
409+
public void testFilterForStructuredPropDeletion() {
410+
Urn deletedPropertyUrn = UrnUtils.getUrn("urn:li:structuredProperty:1");
411+
412+
final CriterionArray criterionArray = new CriterionArray();
413+
criterionArray.add(
414+
new Criterion()
415+
.setField("structuredPropertyPromptUrns")
416+
.setValues(new StringArray(deletedPropertyUrn.toString()))
417+
.setNegated(false)
418+
.setValue("")
419+
.setCondition(Condition.EQUAL));
420+
Filter expectedFilter =
421+
new Filter()
422+
.setOr(
423+
new ConjunctiveCriterionArray(new ConjunctiveCriterion().setAnd(criterionArray)));
424+
425+
assertEquals(
426+
DeleteEntityUtils.getFilterForStructuredPropertyDeletion(deletedPropertyUrn),
427+
expectedFilter);
428+
}
429+
430+
@Test
431+
public void testEntityNamesForStructuredPropDeletion() {
432+
assertEquals(
433+
DeleteEntityUtils.getEntityNamesForStructuredPropertyDeletion(), ImmutableList.of("form"));
434+
}
362435
}

metadata-models/src/main/pegasus/com/linkedin/form/FormPrompt.pdl

+4
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,10 @@ record FormPrompt {
4343
/**
4444
* The structured property that is required on this entity
4545
*/
46+
@Searchable = {
47+
"fieldType": "URN",
48+
"fieldName": "structuredPropertyPromptUrns",
49+
}
4650
urn: Urn
4751
}
4852

metadata-service/factories/src/main/java/com/linkedin/metadata/boot/factories/BootstrapManagerFactory.java

+5-1
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import com.linkedin.metadata.boot.steps.RemoveClientIdAspectStep;
1919
import com.linkedin.metadata.boot.steps.RestoreColumnLineageIndices;
2020
import com.linkedin.metadata.boot.steps.RestoreDbtSiblingsIndices;
21+
import com.linkedin.metadata.boot.steps.RestoreFormInfoIndicesStep;
2122
import com.linkedin.metadata.boot.steps.RestoreGlossaryIndices;
2223
import com.linkedin.metadata.boot.steps.WaitForSystemUpdateStep;
2324
import com.linkedin.metadata.entity.AspectMigrationsDao;
@@ -110,6 +111,8 @@ protected BootstrapManager createInstance(
110111
final WaitForSystemUpdateStep waitForSystemUpdateStep =
111112
new WaitForSystemUpdateStep(_dataHubUpgradeKafkaListener, _configurationProvider);
112113
final IngestEntityTypesStep ingestEntityTypesStep = new IngestEntityTypesStep(_entityService);
114+
final RestoreFormInfoIndicesStep restoreFormInfoIndicesStep =
115+
new RestoreFormInfoIndicesStep(_entityService);
113116

114117
final List<BootstrapStep> finalSteps =
115118
new ArrayList<>(
@@ -124,7 +127,8 @@ protected BootstrapManager createInstance(
124127
restoreDbtSiblingsIndices,
125128
indexDataPlatformsStep,
126129
restoreColumnLineageIndices,
127-
ingestEntityTypesStep));
130+
ingestEntityTypesStep,
131+
restoreFormInfoIndicesStep));
128132

129133
return new BootstrapManager(finalSteps);
130134
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
package com.linkedin.metadata.boot.steps;
2+
3+
import com.linkedin.common.AuditStamp;
4+
import com.linkedin.common.urn.Urn;
5+
import com.linkedin.data.template.RecordTemplate;
6+
import com.linkedin.events.metadata.ChangeType;
7+
import com.linkedin.form.FormInfo;
8+
import com.linkedin.metadata.Constants;
9+
import com.linkedin.metadata.boot.UpgradeStep;
10+
import com.linkedin.metadata.entity.EntityService;
11+
import com.linkedin.metadata.entity.ListResult;
12+
import com.linkedin.metadata.models.AspectSpec;
13+
import com.linkedin.metadata.query.ExtraInfo;
14+
import io.datahubproject.metadata.context.OperationContext;
15+
import java.util.LinkedList;
16+
import java.util.List;
17+
import java.util.Objects;
18+
import java.util.concurrent.ExecutionException;
19+
import java.util.concurrent.Future;
20+
import javax.annotation.Nonnull;
21+
import lombok.extern.slf4j.Slf4j;
22+
23+
@Slf4j
24+
public class RestoreFormInfoIndicesStep extends UpgradeStep {
25+
private static final String VERSION = "1";
26+
private static final String UPGRADE_ID = "restore-form-info-indices";
27+
private static final Integer BATCH_SIZE = 1000;
28+
29+
public RestoreFormInfoIndicesStep(@Nonnull final EntityService<?> entityService) {
30+
super(entityService, VERSION, UPGRADE_ID);
31+
}
32+
33+
@Override
34+
public void upgrade(@Nonnull OperationContext systemOperationContext) throws Exception {
35+
final AuditStamp auditStamp =
36+
new AuditStamp()
37+
.setActor(Urn.createFromString(Constants.SYSTEM_ACTOR))
38+
.setTime(System.currentTimeMillis());
39+
40+
final int totalFormCount = getAndRestoreFormInfoIndices(systemOperationContext, 0, auditStamp);
41+
int formCount = BATCH_SIZE;
42+
while (formCount < totalFormCount) {
43+
getAndRestoreFormInfoIndices(systemOperationContext, formCount, auditStamp);
44+
formCount += BATCH_SIZE;
45+
}
46+
}
47+
48+
@Nonnull
49+
@Override
50+
public ExecutionMode getExecutionMode() {
51+
return ExecutionMode.ASYNC;
52+
}
53+
54+
private int getAndRestoreFormInfoIndices(
55+
@Nonnull OperationContext systemOperationContext, int start, AuditStamp auditStamp) {
56+
final AspectSpec formInfoAspectSpec =
57+
systemOperationContext
58+
.getEntityRegistry()
59+
.getEntitySpec(Constants.FORM_ENTITY_NAME)
60+
.getAspectSpec(Constants.FORM_INFO_ASPECT_NAME);
61+
62+
final ListResult<RecordTemplate> latestAspects =
63+
entityService.listLatestAspects(
64+
systemOperationContext,
65+
Constants.FORM_ENTITY_NAME,
66+
Constants.FORM_INFO_ASPECT_NAME,
67+
start,
68+
BATCH_SIZE);
69+
70+
if (latestAspects.getTotalCount() == 0
71+
|| latestAspects.getValues() == null
72+
|| latestAspects.getMetadata() == null) {
73+
log.debug("Found 0 formInfo aspects for forms. Skipping migration.");
74+
return 0;
75+
}
76+
77+
if (latestAspects.getValues().size() != latestAspects.getMetadata().getExtraInfos().size()) {
78+
// Bad result -- we should log that we cannot migrate this batch of formInfos.
79+
log.warn(
80+
"Failed to match formInfo aspects with corresponding urns. Found mismatched length between aspects ({})"
81+
+ "and metadata ({}) for metadata {}",
82+
latestAspects.getValues().size(),
83+
latestAspects.getMetadata().getExtraInfos().size(),
84+
latestAspects.getMetadata());
85+
return latestAspects.getTotalCount();
86+
}
87+
88+
List<Future<?>> futures = new LinkedList<>();
89+
for (int i = 0; i < latestAspects.getValues().size(); i++) {
90+
ExtraInfo info = latestAspects.getMetadata().getExtraInfos().get(i);
91+
RecordTemplate formInfoRecord = latestAspects.getValues().get(i);
92+
Urn urn = info.getUrn();
93+
FormInfo formInfo = (FormInfo) formInfoRecord;
94+
if (formInfo == null) {
95+
log.warn("Received null formInfo for urn {}", urn);
96+
continue;
97+
}
98+
99+
futures.add(
100+
entityService
101+
.alwaysProduceMCLAsync(
102+
systemOperationContext,
103+
urn,
104+
Constants.FORM_ENTITY_NAME,
105+
Constants.FORM_INFO_ASPECT_NAME,
106+
formInfoAspectSpec,
107+
null,
108+
formInfo,
109+
null,
110+
null,
111+
auditStamp,
112+
ChangeType.RESTATE)
113+
.getFirst());
114+
}
115+
116+
futures.stream()
117+
.filter(Objects::nonNull)
118+
.forEach(
119+
f -> {
120+
try {
121+
f.get();
122+
} catch (InterruptedException | ExecutionException e) {
123+
throw new RuntimeException(e);
124+
}
125+
});
126+
127+
return latestAspects.getTotalCount();
128+
}
129+
}

0 commit comments

Comments
 (0)