feat(restore-indices): createDefaultAspects argument (#12859)

Author: david-leifker

datahub-upgrade/build.gradle

@@ -88,6 +88,9 @@ dependencies {
   testImplementation 'uk.org.webcompere:system-stubs-testng:2.1.7'
   testRuntimeOnly externalDependency.logbackClassic
 
+  testImplementation externalDependency.h2
+  testImplementation testFixtures(project(':metadata-io'))
+
   constraints {
     implementation(implementation externalDependency.parquetHadoop) {
       because("CVE-2022-42003")
@@ -137,7 +140,7 @@ task runRestoreIndices(type: Exec) {
           "-jar",
           "-Dkafka.schemaRegistry.url=http://localhost:8080/schema-registry/api",
           "-Dserver.port=8083",
-          bootJar.getArchiveFile().get(), "-u", "RestoreIndices", "-a", "batchSize=100"
+          bootJar.getArchiveFile().get(), "-u", "RestoreIndices", "-a", "batchSize=100", "-a", "createDefaultAspects=true"
 }
 
 task runRestoreIndicesUrn(type: Exec) {

datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restoreindices/RestoreIndices.java

@@ -26,6 +26,7 @@ public class RestoreIndices implements Upgrade {
   public static final String URN_ARG_NAME = "urn";
   public static final String URN_LIKE_ARG_NAME = "urnLike";
   public static final String URN_BASED_PAGINATION_ARG_NAME = "urnBasedPagination";
+  public static final String CREATE_DEFAULT_ASPECTS_ARG_NAME = "createDefaultAspects";
 
   public static final String STARTING_OFFSET_ARG_NAME = "startingOffset";
   public static final String LAST_URN_ARG_NAME = "lastUrn";

datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restoreindices/SendMAEStep.java

@@ -2,6 +2,7 @@
 
 import static com.linkedin.metadata.Constants.ASPECT_LATEST_VERSION;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.linkedin.datahub.upgrade.UpgradeContext;
 import com.linkedin.datahub.upgrade.UpgradeStep;
 import com.linkedin.datahub.upgrade.UpgradeStepResult;
@@ -80,9 +81,11 @@ private List<RestoreIndicesResult> iterateFutures(List<Future<RestoreIndicesResu
       if (future.isDone()) {
         try {
           result.add(future.get());
-          futures.remove(future);
         } catch (InterruptedException | ExecutionException e) {
           log.error("Error iterating futures", e);
+          result.add(null); // add null to indicate failure
+        } finally {
+          futures.remove(future);
         }
       }
     }
@@ -101,6 +104,7 @@ private RestoreIndicesArgs getArgs(UpgradeContext context) {
     result.batchDelayMs = getBatchDelayMs(context.parsedArgs());
     result.start = getStartingOffset(context.parsedArgs());
     result.urnBasedPagination = getUrnBasedPagination(context.parsedArgs());
+    result.createDefaultAspects = getCreateDefaultAspects(context.parsedArgs());
     if (containsKey(context.parsedArgs(), RestoreIndices.ASPECT_NAME_ARG_NAME)) {
       result.aspectName = context.parsedArgs().get(RestoreIndices.ASPECT_NAME_ARG_NAME).get();
       context.report().addLine(String.format("aspect is %s", result.aspectName));
@@ -151,7 +155,8 @@ private RestoreIndicesArgs getArgs(UpgradeContext context) {
     return result;
   }
 
-  private int getRowCount(RestoreIndicesArgs args) {
+  @VisibleForTesting
+  int getRowCount(RestoreIndicesArgs args) {
     ExpressionList<EbeanAspectV2> countExp =
         _server
             .find(EbeanAspectV2.class)
@@ -233,6 +238,10 @@ public Function<UpgradeContext, UpgradeStepResult> executable() {
         while (futures.size() > 0) {
           List<RestoreIndicesResult> tmpResults = iterateFutures(futures);
           for (RestoreIndicesResult tmpResult : tmpResults) {
+            if (tmpResult == null) {
+              // return error if there was an error processing a future
+              return new DefaultUpgradeStepResult(id(), DataHubUpgradeState.FAILED);
+            }
             reportStats(context, finalJobResult, tmpResult, rowCount, startTime);
           }
         }
@@ -315,6 +324,16 @@ private long getBatchDelayMs(final Map<String, Optional<String>> parsedArgs) {
     return resolvedBatchDelayMs;
   }
 
+  private boolean getCreateDefaultAspects(final Map<String, Optional<String>> parsedArgs) {
+    Boolean createDefaultAspects = null;
+    if (containsKey(parsedArgs, RestoreIndices.CREATE_DEFAULT_ASPECTS_ARG_NAME)) {
+      createDefaultAspects =
+          Boolean.parseBoolean(
+              parsedArgs.get(RestoreIndices.CREATE_DEFAULT_ASPECTS_ARG_NAME).get());
+    }
+    return createDefaultAspects != null ? createDefaultAspects : false;
+  }
+
   private int getThreadCount(final Map<String, Optional<String>> parsedArgs) {
     return getInt(parsedArgs, DEFAULT_THREADS, RestoreIndices.NUM_THREADS_ARG_NAME);
   }

datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/AbstractMCLStep.java

@@ -113,7 +113,9 @@ public Function<UpgradeContext, UpgradeStepResult> executable() {
                   List<Pair<Future<?>, SystemAspect>> futures;
                   futures =
                       EntityUtils.toSystemAspectFromEbeanAspects(
-                              opContext.getRetrieverContext(), batch.collect(Collectors.toList()))
+                              opContext.getRetrieverContext(),
+                              batch.collect(Collectors.toList()),
+                              false)
                           .stream()
                           .map(
                               systemAspect -> {

datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/schemafield/GenerateSchemaFieldsFromSchemaMetadataStep.java

@@ -175,7 +175,8 @@ public Function<UpgradeContext, UpgradeStepResult> executable() {
                                       ebeanAspectV2 ->
                                           EntityUtils.toSystemAspectFromEbeanAspects(
                                               opContext.getRetrieverContext(),
-                                              Set.of(ebeanAspectV2))
+                                              Set.of(ebeanAspectV2),
+                                              true)
                                               .stream())
                                   .map(
                                       systemAspect ->