From bfac935e38a0e0dd19486a26b2691c7034780ac8 Mon Sep 17 00:00:00 2001 From: Masterchen09 <13187726+Masterchen09@users.noreply.github.com> Date: Tue, 4 Jun 2024 20:34:59 +0200 Subject: [PATCH] feat(restore-indices): add additional step to also clear system metadata service --- .../steps/ClearSystemMetadataServiceStep.java | 56 +++++++++++++++++++ .../common/steps/GMSDisableWriteModeStep.java | 4 +- .../common/steps/GMSEnableWriteModeStep.java | 4 +- .../upgrade/config/RestoreBackupConfig.java | 24 ++++++-- .../upgrade/config/RestoreIndicesConfig.java | 16 +++++- .../upgrade/restorebackup/RestoreBackup.java | 31 ++++++---- .../restoreindices/RestoreIndices.java | 9 ++- .../elastic/ElasticSearchGraphService.java | 2 - .../ElasticSearchSystemMetadataService.java | 2 - 9 files changed, 120 insertions(+), 28 deletions(-) create mode 100644 datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/common/steps/ClearSystemMetadataServiceStep.java diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/common/steps/ClearSystemMetadataServiceStep.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/common/steps/ClearSystemMetadataServiceStep.java new file mode 100644 index 00000000000000..889d2f0a79edf6 --- /dev/null +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/common/steps/ClearSystemMetadataServiceStep.java @@ -0,0 +1,56 @@ +package com.linkedin.datahub.upgrade.common.steps; + +import com.linkedin.datahub.upgrade.UpgradeContext; +import com.linkedin.datahub.upgrade.UpgradeStep; +import com.linkedin.datahub.upgrade.UpgradeStepResult; +import com.linkedin.datahub.upgrade.impl.DefaultUpgradeStepResult; +import com.linkedin.datahub.upgrade.nocode.NoCodeUpgrade; +import com.linkedin.metadata.systemmetadata.SystemMetadataService; +import java.util.function.Function; + +public class ClearSystemMetadataServiceStep implements UpgradeStep { + + private final SystemMetadataService _systemMetadataService; + private final boolean _alwaysRun; + + public ClearSystemMetadataServiceStep( + final SystemMetadataService systemMetadataService, final boolean alwaysRun) { + _systemMetadataService = systemMetadataService; + _alwaysRun = alwaysRun; + } + + @Override + public String id() { + return "ClearSystemMetadataServiceStep"; + } + + @Override + public boolean skip(UpgradeContext context) { + if (_alwaysRun) { + return false; + } + if (context.parsedArgs().containsKey(NoCodeUpgrade.CLEAN_ARG_NAME)) { + return false; + } + context.report().addLine("Cleanup has not been requested."); + return true; + } + + @Override + public int retryCount() { + return 1; + } + + @Override + public Function executable() { + return (context) -> { + try { + _systemMetadataService.clear(); + } catch (Exception e) { + context.report().addLine("Failed to clear system metadata service", e); + return new DefaultUpgradeStepResult(id(), UpgradeStepResult.Result.FAILED); + } + return new DefaultUpgradeStepResult(id(), UpgradeStepResult.Result.SUCCEEDED); + }; + } +} diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/common/steps/GMSDisableWriteModeStep.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/common/steps/GMSDisableWriteModeStep.java index 42b3c6b3ccc490..a80adabc60e743 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/common/steps/GMSDisableWriteModeStep.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/common/steps/GMSDisableWriteModeStep.java @@ -13,7 +13,7 @@ @RequiredArgsConstructor public class GMSDisableWriteModeStep implements UpgradeStep { - private final SystemEntityClient entityClient; + private final SystemEntityClient systemEntityClient; @Override public String id() { @@ -29,7 +29,7 @@ public int retryCount() { public Function executable() { return (context) -> { try { - entityClient.setWritable(context.opContext(), false); + systemEntityClient.setWritable(context.opContext(), false); } catch (Exception e) { log.error("Failed to turn write mode off in GMS", e); context.report().addLine("Failed to turn write mode off in GMS"); diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/common/steps/GMSEnableWriteModeStep.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/common/steps/GMSEnableWriteModeStep.java index c43644fda7ddf5..a6b2942fabffd7 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/common/steps/GMSEnableWriteModeStep.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/common/steps/GMSEnableWriteModeStep.java @@ -12,7 +12,7 @@ @Slf4j @RequiredArgsConstructor public class GMSEnableWriteModeStep implements UpgradeStep { - private final SystemEntityClient entityClient; + private final SystemEntityClient systemEntityClient; @Override public String id() { @@ -28,7 +28,7 @@ public int retryCount() { public Function executable() { return (context) -> { try { - entityClient.setWritable(context.opContext(), true); + systemEntityClient.setWritable(context.opContext(), true); } catch (Exception e) { log.error("Failed to turn write mode back on in GMS", e); context.report().addLine("Failed to turn write mode back on in GMS"); diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/RestoreBackupConfig.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/RestoreBackupConfig.java index ec6e5a4a8f04d1..b4eafb4ad3d241 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/RestoreBackupConfig.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/RestoreBackupConfig.java @@ -6,6 +6,7 @@ import com.linkedin.metadata.graph.GraphService; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.search.EntitySearchService; +import com.linkedin.metadata.systemmetadata.SystemMetadataService; import io.ebean.Database; import javax.annotation.Nonnull; import lombok.extern.slf4j.Slf4j; @@ -26,8 +27,9 @@ public class RestoreBackupConfig { "ebeanServer", "entityService", "systemEntityClient", - "graphService", + "systemMetadataService", "searchService", + "graphService", "entityRegistry" }) @ConditionalOnProperty(name = "entityService.impl", havingValue = "ebean", matchIfMissing = true) @@ -35,13 +37,23 @@ public class RestoreBackupConfig { public RestoreBackup createInstance() { final Database ebeanServer = applicationContext.getBean(Database.class); final EntityService entityService = applicationContext.getBean(EntityService.class); - final SystemEntityClient entityClient = applicationContext.getBean(SystemEntityClient.class); - final GraphService graphClient = applicationContext.getBean(GraphService.class); - final EntitySearchService searchClient = applicationContext.getBean(EntitySearchService.class); + final SystemEntityClient systemEntityClient = + applicationContext.getBean(SystemEntityClient.class); + final SystemMetadataService systemMetadataService = + applicationContext.getBean(SystemMetadataService.class); + final EntitySearchService entitySearchService = + applicationContext.getBean(EntitySearchService.class); + final GraphService graphService = applicationContext.getBean(GraphService.class); final EntityRegistry entityRegistry = applicationContext.getBean(EntityRegistry.class); return new RestoreBackup( - ebeanServer, entityService, entityRegistry, entityClient, graphClient, searchClient); + ebeanServer, + entityService, + entityRegistry, + systemEntityClient, + systemMetadataService, + entitySearchService, + graphService); } @Bean(name = "restoreBackup") @@ -49,6 +61,6 @@ public RestoreBackup createInstance() { @Nonnull public RestoreBackup createNotImplInstance() { log.warn("restoreIndices is not supported for cassandra!"); - return new RestoreBackup(null, null, null, null, null, null); + return new RestoreBackup(null, null, null, null, null, null, null); } } diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/RestoreIndicesConfig.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/RestoreIndicesConfig.java index 008bdf5cfac388..949b75edaa6ba0 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/RestoreIndicesConfig.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/RestoreIndicesConfig.java @@ -4,6 +4,7 @@ import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.graph.GraphService; import com.linkedin.metadata.search.EntitySearchService; +import com.linkedin.metadata.systemmetadata.SystemMetadataService; import io.ebean.Database; import javax.annotation.Nonnull; import lombok.extern.slf4j.Slf4j; @@ -20,17 +21,26 @@ public class RestoreIndicesConfig { @Autowired ApplicationContext applicationContext; @Bean(name = "restoreIndices") - @DependsOn({"ebeanServer", "entityService", "searchService", "graphService"}) + @DependsOn({ + "ebeanServer", + "entityService", + "systemMetadataService", + "searchService", + "graphService" + }) @ConditionalOnProperty(name = "entityService.impl", havingValue = "ebean", matchIfMissing = true) @Nonnull public RestoreIndices createInstance() { final Database ebeanServer = applicationContext.getBean(Database.class); final EntityService entityService = applicationContext.getBean(EntityService.class); + final SystemMetadataService systemMetadataService = + applicationContext.getBean(SystemMetadataService.class); final EntitySearchService entitySearchService = applicationContext.getBean(EntitySearchService.class); final GraphService graphService = applicationContext.getBean(GraphService.class); - return new RestoreIndices(ebeanServer, entityService, entitySearchService, graphService); + return new RestoreIndices( + ebeanServer, entityService, systemMetadataService, entitySearchService, graphService); } @Bean(name = "restoreIndices") @@ -38,6 +48,6 @@ public RestoreIndices createInstance() { @Nonnull public RestoreIndices createNotImplInstance() { log.warn("restoreIndices is not supported for cassandra!"); - return new RestoreIndices(null, null, null, null); + return new RestoreIndices(null, null, null, null, null); } } diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restorebackup/RestoreBackup.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restorebackup/RestoreBackup.java index bcaeaa34e8936d..7496655e581b09 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restorebackup/RestoreBackup.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restorebackup/RestoreBackup.java @@ -6,6 +6,7 @@ import com.linkedin.datahub.upgrade.UpgradeStep; import com.linkedin.datahub.upgrade.common.steps.ClearGraphServiceStep; import com.linkedin.datahub.upgrade.common.steps.ClearSearchServiceStep; +import com.linkedin.datahub.upgrade.common.steps.ClearSystemMetadataServiceStep; import com.linkedin.datahub.upgrade.common.steps.GMSDisableWriteModeStep; import com.linkedin.datahub.upgrade.common.steps.GMSEnableWriteModeStep; import com.linkedin.entity.client.SystemEntityClient; @@ -13,6 +14,7 @@ import com.linkedin.metadata.graph.GraphService; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.search.EntitySearchService; +import com.linkedin.metadata.systemmetadata.SystemMetadataService; import io.ebean.Database; import java.util.ArrayList; import java.util.List; @@ -26,13 +28,20 @@ public RestoreBackup( @Nullable final Database server, final EntityService entityService, final EntityRegistry entityRegistry, - final SystemEntityClient entityClient, - final GraphService graphClient, - final EntitySearchService searchClient) { + final SystemEntityClient systemEntityClient, + final SystemMetadataService systemMetadataService, + final EntitySearchService entitySearchService, + final GraphService graphClient) { if (server != null) { _steps = buildSteps( - server, entityService, entityRegistry, entityClient, graphClient, searchClient); + server, + entityService, + entityRegistry, + systemEntityClient, + systemMetadataService, + entitySearchService, + graphClient); } else { _steps = List.of(); } @@ -52,16 +61,18 @@ private List buildSteps( final Database server, final EntityService entityService, final EntityRegistry entityRegistry, - final SystemEntityClient entityClient, - final GraphService graphClient, - final EntitySearchService searchClient) { + final SystemEntityClient systemEntityClient, + final SystemMetadataService systemMetadataService, + final EntitySearchService entitySearchService, + final GraphService graphClient) { final List steps = new ArrayList<>(); - steps.add(new GMSDisableWriteModeStep(entityClient)); - steps.add(new ClearSearchServiceStep(searchClient, true)); + steps.add(new GMSDisableWriteModeStep(systemEntityClient)); + steps.add(new ClearSystemMetadataServiceStep(systemMetadataService, true)); + steps.add(new ClearSearchServiceStep(entitySearchService, true)); steps.add(new ClearGraphServiceStep(graphClient, true)); steps.add(new ClearAspectV2TableStep(server)); steps.add(new RestoreStorageStep(entityService, entityRegistry)); - steps.add(new GMSEnableWriteModeStep(entityClient)); + steps.add(new GMSEnableWriteModeStep(systemEntityClient)); return steps; } diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restoreindices/RestoreIndices.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restoreindices/RestoreIndices.java index 9bc42e23a99746..9d239a56224862 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restoreindices/RestoreIndices.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restoreindices/RestoreIndices.java @@ -6,9 +6,11 @@ import com.linkedin.datahub.upgrade.UpgradeStep; import com.linkedin.datahub.upgrade.common.steps.ClearGraphServiceStep; import com.linkedin.datahub.upgrade.common.steps.ClearSearchServiceStep; +import com.linkedin.datahub.upgrade.common.steps.ClearSystemMetadataServiceStep; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.graph.GraphService; import com.linkedin.metadata.search.EntitySearchService; +import com.linkedin.metadata.systemmetadata.SystemMetadataService; import io.ebean.Database; import java.util.ArrayList; import java.util.List; @@ -32,10 +34,13 @@ public class RestoreIndices implements Upgrade { public RestoreIndices( @Nullable final Database server, final EntityService entityService, + final SystemMetadataService systemMetadataService, final EntitySearchService entitySearchService, final GraphService graphService) { if (server != null) { - _steps = buildSteps(server, entityService, entitySearchService, graphService); + _steps = + buildSteps( + server, entityService, systemMetadataService, entitySearchService, graphService); } else { _steps = List.of(); } @@ -54,9 +59,11 @@ public List steps() { private List buildSteps( final Database server, final EntityService entityService, + final SystemMetadataService systemMetadataService, final EntitySearchService entitySearchService, final GraphService graphService) { final List steps = new ArrayList<>(); + steps.add(new ClearSystemMetadataServiceStep(systemMetadataService, false)); steps.add(new ClearSearchServiceStep(entitySearchService, false)); steps.add(new ClearGraphServiceStep(graphService, false)); steps.add(new SendMAEStep(server, entityService)); diff --git a/metadata-io/src/main/java/com/linkedin/metadata/graph/elastic/ElasticSearchGraphService.java b/metadata-io/src/main/java/com/linkedin/metadata/graph/elastic/ElasticSearchGraphService.java index 37cc6ab6b1eeb3..67523512f7382b 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/graph/elastic/ElasticSearchGraphService.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/graph/elastic/ElasticSearchGraphService.java @@ -4,7 +4,6 @@ import com.fasterxml.jackson.databind.node.JsonNodeFactory; import com.fasterxml.jackson.databind.node.ObjectNode; -import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ImmutableList; import com.linkedin.common.urn.Urn; import com.linkedin.metadata.aspect.models.graph.Edge; @@ -305,7 +304,6 @@ public void reindexAll() { configure(); } - @VisibleForTesting @Override public void clear() { _esBulkProcessor.deleteByQuery( diff --git a/metadata-io/src/main/java/com/linkedin/metadata/systemmetadata/ElasticSearchSystemMetadataService.java b/metadata-io/src/main/java/com/linkedin/metadata/systemmetadata/ElasticSearchSystemMetadataService.java index 36eab7b69e6a12..b0f0597728ee35 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/systemmetadata/ElasticSearchSystemMetadataService.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/systemmetadata/ElasticSearchSystemMetadataService.java @@ -2,7 +2,6 @@ import com.fasterxml.jackson.databind.node.JsonNodeFactory; import com.fasterxml.jackson.databind.node.ObjectNode; -import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ImmutableMap; import com.linkedin.metadata.run.AspectRowSummary; import com.linkedin.metadata.run.IngestionRunSummary; @@ -258,7 +257,6 @@ public void reindexAll() { configure(); } - @VisibleForTesting @Override public void clear() { _esBulkProcessor.deleteByQuery(