fix(restoreIndices): fix bug in urn paginated restoreIndices exit code (datahub-project#11443)

david-leifker · web-flow · commit f6dde2bf0347 · 2024-09-20T10:33:37.000-05:00
diff --git a/datahub-upgrade/build.gradle b/datahub-upgrade/build.gradle
@@ -121,6 +121,34 @@ task run(type: Exec) {
           "-Dserver.port=8083", bootJar.getArchiveFile().get(), "-u", "SystemUpdate"
 }
 
+/**
+ * Runs RestoreIndices on locally running system. The batchSize are set to
+ * test the process with pagination and not designed for optimal performance.
+ */
+task runRestoreIndices(type: Exec) {
+  dependsOn bootJar
+  group = "Execution"
+  description = "Run the restore indices process locally."
+  environment "ENTITY_REGISTRY_CONFIG_PATH", "../metadata-models/src/main/resources/entity-registry.yml"
+  commandLine "java", "-agentlib:jdwp=transport=dt_socket,address=5003,server=y,suspend=n",
+          "-jar",
+          "-Dkafka.schemaRegistry.url=http://localhost:8080/schema-registry/api",
+          "-Dserver.port=8083",
+          bootJar.getArchiveFile().get(), "-u", "RestoreIndices", "-a", "batchSize=100"
+}
+
+task runRestoreIndicesUrn(type: Exec) {
+  dependsOn bootJar
+  group = "Execution"
+  description = "Run the restore indices process locally."
+  environment "ENTITY_REGISTRY_CONFIG_PATH", "../metadata-models/src/main/resources/entity-registry.yml"
+  commandLine "java", "-agentlib:jdwp=transport=dt_socket,address=5003,server=y,suspend=n",
+          "-jar",
+          "-Dkafka.schemaRegistry.url=http://localhost:8080/schema-registry/api",
+          "-Dserver.port=8083",
+          bootJar.getArchiveFile().get(), "-u", "RestoreIndices", "-a", "batchSize=100", "-a", "urnBasedPagination=true"
+}
+
 docker {
   name "${docker_registry}/${docker_repo}:v${version}"
   version "v${version}"
diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restoreindices/SendMAEStep.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restoreindices/SendMAEStep.java
@@ -16,6 +16,7 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+import java.util.NoSuchElementException;
 import java.util.Optional;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ExecutionException;
@@ -189,7 +190,12 @@ public Function<UpgradeContext, UpgradeStepResult> executable() {
             context.report().addLine(String.format("Rows processed this loop %d", rowsProcessed));
             start += args.batchSize;
           } catch (InterruptedException | ExecutionException e) {
-            return new DefaultUpgradeStepResult(id(), DataHubUpgradeState.FAILED);
+            if (e.getCause() instanceof NoSuchElementException) {
+              context.report().addLine("End of data.");
+              break;
+            } else {
+              return new DefaultUpgradeStepResult(id(), DataHubUpgradeState.FAILED);
+            }
           }
         }
       } else {
