diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/UpgradeCli.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/UpgradeCli.java index e7c730d6e6d249..98e8de357f2322 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/UpgradeCli.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/UpgradeCli.java @@ -3,6 +3,7 @@ import com.linkedin.datahub.upgrade.impl.DefaultUpgradeManager; import com.linkedin.datahub.upgrade.nocode.NoCodeUpgrade; import com.linkedin.datahub.upgrade.nocodecleanup.NoCodeCleanupUpgrade; +import com.linkedin.datahub.upgrade.removeunknownaspects.RemoveUnknownAspects; import com.linkedin.datahub.upgrade.restorebackup.RestoreBackup; import com.linkedin.datahub.upgrade.restoreindices.RestoreIndices; import java.util.List; @@ -44,12 +45,17 @@ private static final class Args { @Named("restoreBackup") private RestoreBackup restoreBackup; + @Inject + @Named("removeUnknownAspects") + private RemoveUnknownAspects removeUnknownAspects; + @Override public void run(String... cmdLineArgs) { _upgradeManager.register(noCodeUpgrade); _upgradeManager.register(noCodeCleanup); _upgradeManager.register(restoreIndices); _upgradeManager.register(restoreBackup); + _upgradeManager.register(removeUnknownAspects); final Args args = new Args(); new CommandLine(args).setCaseInsensitiveEnumValuesAllowed(true).parseArgs(cmdLineArgs); diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/RemoveUnknownAspectsConfig.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/RemoveUnknownAspectsConfig.java new file mode 100644 index 00000000000000..cdc739efc416dd --- /dev/null +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/RemoveUnknownAspectsConfig.java @@ -0,0 +1,15 @@ +package com.linkedin.datahub.upgrade.config; + +import com.linkedin.datahub.upgrade.removeunknownaspects.RemoveUnknownAspects; +import com.linkedin.metadata.entity.EntityService; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; + + +@Configuration +public class RemoveUnknownAspectsConfig { + @Bean(name = "removeUnknownAspects") + public RemoveUnknownAspects removeUnknownAspects(EntityService entityService) { + return new RemoveUnknownAspects(entityService); + } +} diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/removeunknownaspects/RemoveClientIdAspectStep.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/removeunknownaspects/RemoveClientIdAspectStep.java new file mode 100644 index 00000000000000..b55d439745e691 --- /dev/null +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/removeunknownaspects/RemoveClientIdAspectStep.java @@ -0,0 +1,41 @@ +package com.linkedin.datahub.upgrade.removeunknownaspects; + +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.gms.factory.telemetry.TelemetryUtils; +import com.linkedin.metadata.entity.EntityService; +import java.util.HashMap; +import java.util.function.Function; +import lombok.RequiredArgsConstructor; +import lombok.extern.slf4j.Slf4j; + + +@Slf4j +@RequiredArgsConstructor +public class RemoveClientIdAspectStep implements UpgradeStep { + + private static final String INVALID_CLIENT_ID_ASPECT = "clientId"; + + private final EntityService _entityService; + + @Override + public String id() { + return this.getClass().getSimpleName(); + } + + @Override + public boolean skip(UpgradeContext context) { + return false; + } + + @Override + public Function executable() { + return upgradeContext -> { + _entityService.deleteAspect(TelemetryUtils.CLIENT_ID_URN, INVALID_CLIENT_ID_ASPECT, + new HashMap<>(), true); + return (UpgradeStepResult) new DefaultUpgradeStepResult(id(), UpgradeStepResult.Result.SUCCEEDED); + }; + } +} diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/removeunknownaspects/RemoveUnknownAspects.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/removeunknownaspects/RemoveUnknownAspects.java new file mode 100644 index 00000000000000..f8af69dba08653 --- /dev/null +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/removeunknownaspects/RemoveUnknownAspects.java @@ -0,0 +1,40 @@ +package com.linkedin.datahub.upgrade.removeunknownaspects; + +import com.google.common.collect.ImmutableList; +import com.linkedin.datahub.upgrade.Upgrade; +import com.linkedin.datahub.upgrade.UpgradeCleanupStep; +import com.linkedin.datahub.upgrade.UpgradeStep; +import com.linkedin.metadata.entity.EntityService; +import java.util.ArrayList; +import java.util.List; + + +public class RemoveUnknownAspects implements Upgrade { + + private final List _steps; + + public RemoveUnknownAspects(final EntityService entityService) { + _steps = buildSteps(entityService); + } + + @Override + public String id() { + return this.getClass().getSimpleName(); + } + + @Override + public List steps() { + return _steps; + } + + private List buildSteps(final EntityService entityService) { + final List steps = new ArrayList<>(); + steps.add(new RemoveClientIdAspectStep(entityService)); + return steps; + } + + @Override + public List cleanupSteps() { + return ImmutableList.of(); + } +} diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityAspectIdentifier.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityAspectIdentifier.java index 00919335516711..cb360192c01201 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityAspectIdentifier.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityAspectIdentifier.java @@ -1,16 +1,28 @@ package com.linkedin.metadata.entity; +import com.linkedin.metadata.entity.cassandra.CassandraAspect; +import com.linkedin.metadata.entity.ebean.EbeanAspectV2; +import javax.annotation.Nonnull; import lombok.Value; +import lombok.extern.slf4j.Slf4j; -import javax.annotation.Nonnull; /** * This class holds values required to construct a unique key to identify an entity aspect record in a database. * Its existence started mainly for compatibility with {@link com.linkedin.metadata.entity.ebean.EbeanAspectV2.PrimaryKey} */ @Value +@Slf4j public class EntityAspectIdentifier { @Nonnull String urn; @Nonnull String aspect; long version; + + public static EntityAspectIdentifier fromEbean(EbeanAspectV2 ebeanAspectV2) { + return new EntityAspectIdentifier(ebeanAspectV2.getUrn(), ebeanAspectV2.getAspect(), ebeanAspectV2.getVersion()); + } + + public static EntityAspectIdentifier fromCassandra(CassandraAspect cassandraAspect) { + return new EntityAspectIdentifier(cassandraAspect.getUrn(), cassandraAspect.getAspect(), cassandraAspect.getVersion()); + } } diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java index a4b61d036c6ef1..5fd7c428750e73 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java @@ -3,7 +3,6 @@ import com.codahale.metrics.Timer; import com.datahub.util.RecordUtils; import com.datahub.util.exception.ModelConversionException; -import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; import com.google.common.collect.Iterators; @@ -44,15 +43,10 @@ import com.linkedin.metadata.utils.PegasusUtils; import com.linkedin.metadata.utils.metrics.MetricUtils; import com.linkedin.mxe.MetadataAuditOperation; -import com.linkedin.mxe.MetadataChangeProposal; import com.linkedin.mxe.MetadataChangeLog; +import com.linkedin.mxe.MetadataChangeProposal; import com.linkedin.mxe.SystemMetadata; import com.linkedin.util.Pair; -import lombok.Value; -import lombok.extern.slf4j.Slf4j; - -import javax.annotation.Nonnull; -import javax.annotation.Nullable; import java.net.URISyntaxException; import java.sql.Timestamp; import java.util.ArrayList; @@ -67,11 +61,13 @@ import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Function; import java.util.stream.Collectors; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import lombok.Value; +import lombok.extern.slf4j.Slf4j; -import static com.linkedin.metadata.Constants.ASPECT_LATEST_VERSION; -import static com.linkedin.metadata.Constants.SYSTEM_ACTOR; -import static com.linkedin.metadata.utils.PegasusUtils.getDataTemplateClassFromSchema; -import static com.linkedin.metadata.utils.PegasusUtils.urnToEntityName; +import static com.linkedin.metadata.Constants.*; +import static com.linkedin.metadata.utils.PegasusUtils.*; /** @@ -1315,7 +1311,7 @@ public Boolean exists(Urn urn) { } @Nullable - public RollbackResult deleteAspect(String urn, String aspectName, Map conditions, boolean hardDelete) { + public RollbackResult deleteAspect(String urn, String aspectName, @Nonnull Map conditions, boolean hardDelete) { // Validate pre-conditions before running queries Urn entityUrn; EntitySpec entitySpec; @@ -1323,8 +1319,6 @@ public RollbackResult deleteAspect(String urn, String aspectName, Map(), true); + + final AuditStamp auditStamp = new AuditStamp().setActor(Urn.createFromString(Constants.SYSTEM_ACTOR)).setTime(System.currentTimeMillis()); + final DataHubUpgradeResult upgradeResult = new DataHubUpgradeResult().setTimestampMs(System.currentTimeMillis()); + ingestUpgradeAspect(Constants.DATA_HUB_UPGRADE_RESULT_ASPECT_NAME, upgradeResult, auditStamp); + } catch (Exception e) { + log.error("Error when running the RemoveUnknownAspects Bootstrap Step", e); + _entityService.deleteUrn(REMOVE_UNKNOWN_ASPECTS_URN); + throw new RuntimeException("Error when running the RemoveUnknownAspects Bootstrap Step", e); + } + } + + @Nonnull + @Override + public ExecutionMode getExecutionMode() { + return ExecutionMode.ASYNC; + } + + private void ingestUpgradeAspect(String aspectName, RecordTemplate aspect, AuditStamp auditStamp) { + final MetadataChangeProposal upgradeProposal = new MetadataChangeProposal(); + upgradeProposal.setEntityUrn(REMOVE_UNKNOWN_ASPECTS_URN); + upgradeProposal.setEntityType(Constants.DATA_HUB_UPGRADE_ENTITY_NAME); + upgradeProposal.setAspectName(aspectName); + upgradeProposal.setAspect(GenericRecordUtils.serializeAspect(aspect)); + upgradeProposal.setChangeType(ChangeType.UPSERT); + + _entityService.ingestProposal(upgradeProposal, auditStamp); + } + +}