feat(config): Configurable bootstrap of ownership types (#9308)

metadata-service/configuration/src/main/resources/application.yml

@@ -292,6 +292,8 @@ bootstrap:
     file: ${BOOTSTRAP_POLICIES_FILE:classpath:boot/policies.json}
     # eg for local file
     # file: "file:///datahub/datahub-gms/resources/custom-policies.json"
+  ownershipTypes:
+    file: ${BOOTSTRAP_OWNERSHIP_TYPES_FILE:classpath:boot/ownership_types.json}
   servlets:
     waitTimeout: ${BOOTSTRAP_SERVLETS_WAITTIMEOUT:60} # Total waiting time in seconds for servlets to initialize
 

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

@@ -94,6 +94,9 @@ public class BootstrapManagerFactory {
   @Value("${bootstrap.policies.file}")
   private Resource _policiesResource;
 
+  @Value("${bootstrap.ownershipTypes.file}")
+  private Resource _ownershipTypesResource;
+
   @Bean(name = "bootstrapManager")
   @Scope("singleton")
   @Nonnull
@@ -116,7 +119,7 @@ protected BootstrapManager createInstance() {
     final IngestDefaultGlobalSettingsStep ingestSettingsStep = new IngestDefaultGlobalSettingsStep(_entityService);
     final WaitForSystemUpdateStep waitForSystemUpdateStep = new WaitForSystemUpdateStep(_dataHubUpgradeKafkaListener,
         _configurationProvider);
-    final IngestOwnershipTypesStep ingestOwnershipTypesStep = new IngestOwnershipTypesStep(_entityService);
+    final IngestOwnershipTypesStep ingestOwnershipTypesStep = new IngestOwnershipTypesStep(_entityService, _ownershipTypesResource);
 
     final List<BootstrapStep> finalSteps = new ArrayList<>(ImmutableList.of(
         waitForSystemUpdateStep,

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

@@ -7,7 +7,7 @@
 import com.linkedin.common.urn.Urn;
 import com.linkedin.events.metadata.ChangeType;
 import com.linkedin.metadata.Constants;
-import com.linkedin.metadata.boot.UpgradeStep;
+import com.linkedin.metadata.boot.BootstrapStep;
 import com.linkedin.metadata.entity.EntityService;
 import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl;
 import com.linkedin.metadata.models.AspectSpec;
@@ -16,9 +16,10 @@
 import com.linkedin.mxe.GenericAspect;
 import com.linkedin.mxe.MetadataChangeProposal;
 import com.linkedin.ownership.OwnershipTypeInfo;
-import javax.annotation.Nonnull;
+
+import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
-import org.springframework.core.io.ClassPathResource;
+import org.springframework.core.io.Resource;
 
 import java.util.List;
 
@@ -33,42 +34,30 @@
  * If not, it ingests the ownership type into DataHub.
  */
 @Slf4j
-public class IngestOwnershipTypesStep extends UpgradeStep {
-
-  private static final String UPGRADE_ID = "ingest-default-metadata-ownership-types";
-  private static final String VERSION = "1";
-  private static final int SLEEP_SECONDS = 60;
+@RequiredArgsConstructor
+public class IngestOwnershipTypesStep implements BootstrapStep {
 
   private static final ObjectMapper JSON_MAPPER = new ObjectMapper();
-
-  public IngestOwnershipTypesStep(EntityService entityService) {
-    super(entityService, VERSION, UPGRADE_ID);
-  }
+  private final EntityService _entityService;
+  private final Resource _ownershipTypesResource;
 
   @Override
   public String name() {
     return "IngestOwnershipTypesStep";
   }
 
   @Override
-  public void upgrade() throws Exception {
-    log.info("Ingesting default ownership types...");
-
-    // Sleep to ensure deployment process finishes.
-    Thread.sleep(SLEEP_SECONDS * 1000);
+  public void execute() throws Exception {
+    log.info("Ingesting default ownership types from {}...", _ownershipTypesResource);
 
     // 1. Read from the file into JSON.
-    final JsonNode ownershipTypesObj = JSON_MAPPER.readTree(new ClassPathResource("./boot/ownership_types.json")
-        .getFile());
+    final JsonNode ownershipTypesObj = JSON_MAPPER.readTree(_ownershipTypesResource.getFile());
 
     if (!ownershipTypesObj.isArray()) {
       throw new RuntimeException(String.format("Found malformed ownership file, expected an Array but found %s",
           ownershipTypesObj.getNodeType()));
     }
 
-    final AspectSpec ownershipTypeInfoAspectSpec = _entityService.getEntityRegistry()
-        .getEntitySpec(OWNERSHIP_TYPE_ENTITY_NAME)
-        .getAspectSpec(OWNERSHIP_TYPE_INFO_ASPECT_NAME);
     final AuditStamp auditStamp =
         new AuditStamp().setActor(Urn.createFromString(Constants.SYSTEM_ACTOR)).setTime(System.currentTimeMillis());
 
@@ -79,14 +68,13 @@ public void upgrade() throws Exception {
       final OwnershipTypeInfo info = RecordUtils.toRecordTemplate(OwnershipTypeInfo.class, roleObj.get("info")
           .toString());
       log.info(String.format("Ingesting default ownership type with urn %s", urn));
-      ingestOwnershipType(urn, info, auditStamp, ownershipTypeInfoAspectSpec);
+      ingestOwnershipType(urn, info, auditStamp);
       numIngested++;
     }
     log.info("Ingested {} new ownership types", numIngested);
   }
 
-  private void ingestOwnershipType(final Urn ownershipTypeUrn, final OwnershipTypeInfo info, final AuditStamp auditStamp,
-      final AspectSpec ownershipTypeInfoAspectSpec) {
+  private void ingestOwnershipType(final Urn ownershipTypeUrn, final OwnershipTypeInfo info, final AuditStamp auditStamp) {
 
     // 3. Write key & aspect MCPs.
     final MetadataChangeProposal keyAspectProposal = new MetadataChangeProposal();
@@ -112,10 +100,4 @@ private void ingestOwnershipType(final Urn ownershipTypeUrn, final OwnershipType
             .mcps(List.of(keyAspectProposal, proposal), _entityService.getEntityRegistry()).build(), auditStamp,
             false);
   }
-
-  @Nonnull
-  @Override
-  public ExecutionMode getExecutionMode() {
-    return ExecutionMode.ASYNC;
-  }
 }