From 9ef281acf2ef0bca23bb3486e47a9a2342f25ff9 Mon Sep 17 00:00:00 2001
From: david-leifker <114954101+david-leifker@users.noreply.github.com>
Date: Sun, 30 Apr 2023 08:12:11 -0500
Subject: [PATCH] fix(system-update): fix no wait flag (#7927)

---
 .../factories/BootstrapManagerFactory.java    | 39 ++++++++++---------
 .../kafka/DataHubUpgradeKafkaListener.java    |  3 ++
 2 files changed, 24 insertions(+), 18 deletions(-)

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 8a177b08b4..dd74c764c4 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
@@ -1,6 +1,5 @@
 package com.linkedin.metadata.boot.factories;
 
-import com.google.common.collect.ImmutableList;
 import com.linkedin.gms.factory.config.ConfigurationProvider;
 import com.linkedin.gms.factory.entity.EntityServiceFactory;
 import com.linkedin.gms.factory.entityregistry.EntityRegistryFactory;
@@ -28,8 +27,10 @@ import com.linkedin.metadata.entity.EntityService;
 import com.linkedin.metadata.models.registry.EntityRegistry;
 import com.linkedin.metadata.search.EntitySearchService;
 import com.linkedin.metadata.search.transformer.SearchDocumentTransformer;
-import java.util.ArrayList;
 import java.util.List;
+import java.util.Objects;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 import javax.annotation.Nonnull;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Qualifier;
@@ -99,23 +100,25 @@ public class BootstrapManagerFactory {
     final RemoveClientIdAspectStep removeClientIdAspectStep = new RemoveClientIdAspectStep(_entityService);
     final RestoreColumnLineageIndices restoreColumnLineageIndices = new RestoreColumnLineageIndices(_entityService, _entityRegistry);
     final IngestDefaultGlobalSettingsStep ingestSettingsStep = new IngestDefaultGlobalSettingsStep(_entityService);
-    final WaitForSystemUpdateStep waitForSystemUpdateStep = new WaitForSystemUpdateStep(_dataHubUpgradeKafkaListener,
-        _configurationProvider);
+    final WaitForSystemUpdateStep waitForSystemUpdateStep = _configurationProvider.getSystemUpdate().isWaitForSystemUpdate()
+        ? new WaitForSystemUpdateStep(_dataHubUpgradeKafkaListener, _configurationProvider) : null;
 
-    final List<BootstrapStep> finalSteps = new ArrayList<>(ImmutableList.of(
-            waitForSystemUpdateStep,
-            ingestRootUserStep,
-            ingestPoliciesStep,
-            ingestRolesStep,
-            ingestDataPlatformsStep,
-            ingestDataPlatformInstancesStep,
-            _ingestRetentionPoliciesStep,
-            ingestSettingsStep,
-            restoreGlossaryIndicesStep,
-            removeClientIdAspectStep,
-            restoreDbtSiblingsIndices,
-            indexDataPlatformsStep,
-            restoreColumnLineageIndices));
+    final List<BootstrapStep> finalSteps = Stream.of(
+                    waitForSystemUpdateStep,
+                    ingestRootUserStep,
+                    ingestPoliciesStep,
+                    ingestRolesStep,
+                    ingestDataPlatformsStep,
+                    ingestDataPlatformInstancesStep,
+                    _ingestRetentionPoliciesStep,
+                    ingestSettingsStep,
+                    restoreGlossaryIndicesStep,
+                    removeClientIdAspectStep,
+                    restoreDbtSiblingsIndices,
+                    indexDataPlatformsStep,
+                    restoreColumnLineageIndices)
+            .filter(Objects::nonNull)
+            .collect(Collectors.toList());
 
     if (_upgradeDefaultBrowsePathsEnabled) {
       finalSteps.add(new UpgradeDefaultBrowsePathsStep(_entityService));
diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/kafka/DataHubUpgradeKafkaListener.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/kafka/DataHubUpgradeKafkaListener.java
index 33a517eede..8afc229c47 100644
--- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/kafka/DataHubUpgradeKafkaListener.java
+++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/kafka/DataHubUpgradeKafkaListener.java
@@ -84,6 +84,9 @@ public class DataHubUpgradeKafkaListener implements ConsumerSeekAware, Bootstrap
       log.info("Latest system update version: {}", event.getVersion());
       if (expectedVersion.equals(event.getVersion())) {
         IS_UPDATED.getAndSet(true);
+      } else if (!_configurationProvider.getSystemUpdate().isWaitForSystemUpdate()) {
+        log.warn("Wait for system update is disabled. Proceeding with startup.");
+        IS_UPDATED.getAndSet(true);
       } else {
         log.warn("System version is not up to date: {}. Waiting for datahub-upgrade to complete...", expectedVersion);
       }