feat(retention): Add retention to Local DB (datahub-project#3715)

Author: Dexter Lee

docs-website/sidebars.js

@@ -228,6 +228,7 @@ module.exports = {
       {
         Advanced: [
           "docs/advanced/no-code-modeling",
+          "docs/advanced/db-retention",
           "docs/advanced/aspect-versioning",
           "docs/advanced/es-7-upgrade",
           "docs/advanced/high-cardinality",

docs/advanced/db-retention.md

@@ -0,0 +1,79 @@
+# Configuring Database Retention
+
+## Goal
+
+DataHub uses a database (or key-value store) to store different versions of the aspects as they get ingested. Storing
+multiple versions of the aspects allows us to look at the history of how the aspect changed and to rollback to previous
+version when incorrect metadata gets ingested. However, each version takes up space in the database, while bringing less
+value to the system. We need to be able to impose **retention** on these records to keep the size of the DB in check.
+
+Goal of the retention system is to be able to **configure and enforce retention policies** on documents in various
+levels (
+global, entity-level, aspect-level)
+
+## What type of retention policies are supported?
+
+We support 3 types of retention policies.
+
+1. Indefinite retention: Keep all versions of aspects
+2. Version-based retention: Keep the latest N versions
+3. Time-based retention: Keep versions that have been ingested in the last N seconds
+
+Note, the latest version of each aspect (version 0) is never deleted. This is to ensure that we do not impact the core
+functionality of DataHub while applying retention.
+
+## When is the retention policy applied?
+
+As of now, retention policies are applied in two places
+
+1. **GMS boot-up**: On boot, it runs a bootstrap step to ingest the predefined set of retention policies. If there were
+   no existing policies or the existing policies got updated, it will trigger an asynchronous call to apply retention
+   to **
+   all** records in the database.
+2. **Ingest**: On every ingest, if an existing aspect got updated, it applies retention to the urn, aspect pair being
+   ingested.
+
+We are planning to support a cron-based application of retention in the near future to ensure that the time-based
+retention is applied correctly.
+
+## How to configure?
+
+For the initial iteration, we have made this feature opt-in. Please set **ENTITY_SERVICE_ENABLE_RETENTION=true** when
+creating the datahub-gms container/k8s pod.
+
+On GMS start up, it fetches the list of retention policies to ingest from two sources. First is the default we provide,
+which adds a version-based retention to keep 20 latest aspects for all entity-aspect pairs. Second, we read YAML files
+from the `/etc/datahub/plugins/retention` directory and overlay them on the default set of policies we provide.
+
+For docker, we set docker-compose to mount `${HOME}/.datahub/plugins` directory to `/etc/datahub/plugins` directory
+within the containers, so you can customize the initial set of retention policies by creating
+a `${HOME}/.datahub/plugins/retention/retention.yaml` file.
+
+We will support a standardized way to do this in kubernetes setup in the near future. 
+
+The format for the YAML file is as follows. 
+
+```yaml
+- entity: "*" # denotes that policy will be applied to all entities
+  aspect: "*" # denotes that policy will be applied to all aspects
+  config:
+    retention:
+      version:
+        maxVersions: 20
+- entity: "dataset"
+  aspect: "datasetProperties"
+  config:
+    retention:
+      version:
+        maxVersions: 20
+      time:
+        maxAgeInSeconds: 2592000 # 30 days
+```
+
+Note, it searches for the policies corresponding to the entity, aspect pair in the following order
+1. entity, aspect
+2. *, aspect
+3. entity, *
+4. *, *
+
+By restarting datahub-gms after creating the plugin yaml file, the new set of retention policies will be applied. 

metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java

@@ -1,5 +1,6 @@
 package com.linkedin.metadata.entity;
 
+import com.codahale.metrics.Timer;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Streams;
 import com.linkedin.common.AuditStamp;
@@ -26,6 +27,7 @@
 import com.linkedin.metadata.utils.EntityKeyUtils;
 import com.linkedin.metadata.utils.GenericAspectUtils;
 import com.linkedin.metadata.utils.PegasusUtils;
+import com.linkedin.metadata.utils.metrics.MetricUtils;
 import com.linkedin.mxe.MetadataAuditOperation;
 import com.linkedin.mxe.MetadataChangeLog;
 import com.linkedin.mxe.MetadataChangeProposal;
@@ -38,9 +40,13 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
+import lombok.Getter;
+import lombok.Setter;
+import lombok.Value;
 import lombok.extern.slf4j.Slf4j;
 
 import static com.linkedin.metadata.Constants.ASPECT_LATEST_VERSION;
@@ -88,7 +94,10 @@ public abstract class EntityService {
   private final EntityEventProducer _producer;
   private final EntityRegistry _entityRegistry;
   private final Map<String, Set<String>> _entityToValidAspects;
-  private Boolean _emitAspectSpecificAuditEvent = false;
+  @Getter
+  @Setter
+  private RetentionService retentionService;
+  private Boolean _alwaysEmitAuditEvent = false;
   public static final String DEFAULT_RUN_ID = "no-run-id-provided";
   public static final String BROWSE_PATHS = "browsePaths";
   public static final String DATA_PLATFORM_INSTANCE = "dataPlatformInstance";
@@ -107,7 +116,7 @@ protected EntityService(@Nonnull final EntityEventProducer producer, @Nonnull fi
    * @param aspectNames aspects to fetch for each urn in urns set
    * @return a map of provided {@link Urn} to a List containing the requested aspects.
    */
-  protected abstract Map<Urn, List<RecordTemplate>> getLatestAspects(@Nonnull final Set<Urn> urns,
+  public abstract Map<Urn, List<RecordTemplate>> getLatestAspects(@Nonnull final Set<Urn> urns,
       @Nonnull final Set<String> aspectNames);
 
   /**
@@ -146,6 +155,22 @@ public abstract VersionedAspect getVersionedAspect(@Nonnull final Urn urn, @Nonn
   public abstract ListResult<RecordTemplate> listLatestAspects(@Nonnull final String entityName,
       @Nonnull final String aspectName, final int start, int count);
 
+  /**
+   * Checks whether there is an actual update to the aspect by applying the updateLambda
+   * If there is an update, push the new version into the local DB.
+   * Otherwise, do not push the new version, but just update the system metadata.
+   *
+   * @param urn an urn associated with the new aspect
+   * @param aspectName name of the aspect being inserted
+   * @param updateLambda Function to apply to the latest version of the aspect to get the updated version
+   * @param auditStamp an {@link AuditStamp} containing metadata about the writer & current time   * @param providedSystemMetadata
+   * @return Details about the new and old version of the aspect
+   */
+  @Nonnull
+  protected abstract UpdateAspectResult ingestAspectToLocalDB(@Nonnull final Urn urn, @Nonnull final String aspectName,
+      @Nonnull final Function<Optional<RecordTemplate>, RecordTemplate> updateLambda,
+      @Nonnull final AuditStamp auditStamp, @Nonnull final SystemMetadata providedSystemMetadata);
+
   /**
    * Ingests (inserts) a new version of an entity aspect & emits a {@link com.linkedin.mxe.MetadataAuditEvent}.
    *
@@ -159,8 +184,47 @@ public abstract ListResult<RecordTemplate> listLatestAspects(@Nonnull final Stri
    * @param systemMetadata
    * @return the {@link RecordTemplate} representation of the written aspect object
    */
-  public abstract RecordTemplate ingestAspect(@Nonnull final Urn urn, @Nonnull final String aspectName,
-      @Nonnull final RecordTemplate newValue, @Nonnull final AuditStamp auditStamp, SystemMetadata systemMetadata);
+  public RecordTemplate ingestAspect(@Nonnull final Urn urn, @Nonnull final String aspectName,
+      @Nonnull final RecordTemplate newValue, @Nonnull final AuditStamp auditStamp, SystemMetadata systemMetadata) {
+
+    log.debug("Invoked ingestAspect with urn: {}, aspectName: {}, newValue: {}", urn, aspectName, newValue);
+
+    if (!urn.toString().trim().equals(urn.toString())) {
+      throw new IllegalArgumentException("Error: cannot provide an URN with leading or trailing whitespace");
+    }
+
+    Timer.Context ingestToLocalDBTimer = MetricUtils.timer(this.getClass(), "ingestAspectToLocalDB").time();
+    UpdateAspectResult result = ingestAspectToLocalDB(urn, aspectName, ignored -> newValue, auditStamp, systemMetadata);
+    ingestToLocalDBTimer.stop();
+
+    final RecordTemplate oldValue = result.getOldValue();
+    final RecordTemplate updatedValue = result.getNewValue();
+
+    // Apply retention policies asynchronously if there was an update to existing aspect value
+    if (oldValue != updatedValue && oldValue != null && retentionService != null) {
+      retentionService.applyRetention(urn, aspectName,
+          Optional.of(new RetentionService.RetentionContext(Optional.of(result.maxVersion))));
+    }
+
+    // Produce MAE after a successful update
+    if (oldValue != updatedValue || _alwaysEmitAuditEvent) {
+      log.debug(String.format("Producing MetadataAuditEvent for ingested aspect %s, urn %s", aspectName, urn));
+      Timer.Context produceMAETimer = MetricUtils.timer(this.getClass(), "produceMAE").time();
+      if (aspectName.equals(getKeyAspectName(urn))) {
+        produceMetadataAuditEventForKey(urn, result.getNewSystemMetadata());
+      } else {
+        produceMetadataAuditEvent(urn, oldValue, updatedValue, result.getOldSystemMetadata(),
+            result.getNewSystemMetadata(), MetadataAuditOperation.UPDATE);
+      }
+      produceMAETimer.stop();
+    } else {
+      log.debug(
+          String.format("Skipped producing MetadataAuditEvent for ingested aspect %s, urn %s. Aspect has not changed.",
+              aspectName, urn));
+    }
+
+    return updatedValue;
+  }
 
   public RecordTemplate ingestAspect(@Nonnull final Urn urn, @Nonnull final String aspectName,
       @Nonnull final RecordTemplate newValue, @Nonnull final AuditStamp auditStamp) {
@@ -171,6 +235,105 @@ public RecordTemplate ingestAspect(@Nonnull final Urn urn, @Nonnull final String
     return ingestAspect(urn, aspectName, newValue, auditStamp, generatedSystemMetadata);
   }
 
+  public IngestProposalResult ingestProposal(@Nonnull MetadataChangeProposal metadataChangeProposal, AuditStamp auditStamp) {
+
+    log.debug("entity type = {}", metadataChangeProposal.getEntityType());
+    EntitySpec entitySpec = getEntityRegistry().getEntitySpec(metadataChangeProposal.getEntityType());
+    log.debug("entity spec = {}", entitySpec);
+
+    Urn entityUrn = EntityKeyUtils.getUrnFromProposal(metadataChangeProposal, entitySpec.getKeyAspectSpec());
+
+    if (metadataChangeProposal.getChangeType() != ChangeType.UPSERT) {
+      throw new UnsupportedOperationException("Only upsert operation is supported");
+    }
+
+    if (!metadataChangeProposal.hasAspectName() || !metadataChangeProposal.hasAspect()) {
+      throw new UnsupportedOperationException("Aspect and aspect name is required for create and update operations");
+    }
+
+    AspectSpec aspectSpec = entitySpec.getAspectSpec(metadataChangeProposal.getAspectName());
+
+    if (aspectSpec == null) {
+      throw new RuntimeException(
+          String.format("Unknown aspect %s for entity %s", metadataChangeProposal.getAspectName(),
+              metadataChangeProposal.getEntityType()));
+    }
+
+    log.debug("aspect spec = {}", aspectSpec);
+
+    RecordTemplate aspect;
+    try {
+      aspect = GenericAspectUtils.deserializeAspect(metadataChangeProposal.getAspect().getValue(),
+          metadataChangeProposal.getAspect().getContentType(), aspectSpec);
+      ValidationUtils.validateOrThrow(aspect);
+    } catch (ModelConversionException e) {
+      throw new RuntimeException(
+          String.format("Could not deserialize {} for aspect {}", metadataChangeProposal.getAspect().getValue(),
+              metadataChangeProposal.getAspectName()));
+    }
+    log.debug("aspect = {}", aspect);
+
+    SystemMetadata systemMetadata = metadataChangeProposal.getSystemMetadata();
+    if (systemMetadata == null) {
+      systemMetadata = new SystemMetadata();
+      systemMetadata.setRunId(DEFAULT_RUN_ID);
+      systemMetadata.setLastObserved(System.currentTimeMillis());
+    }
+    systemMetadata.setRegistryName(aspectSpec.getRegistryName());
+    systemMetadata.setRegistryVersion(aspectSpec.getRegistryVersion().toString());
+
+    RecordTemplate oldAspect = null;
+    SystemMetadata oldSystemMetadata = null;
+    RecordTemplate newAspect = aspect;
+    SystemMetadata newSystemMetadata = systemMetadata;
+
+    if (!aspectSpec.isTimeseries()) {
+      Timer.Context ingestToLocalDBTimer = MetricUtils.timer(this.getClass(), "ingestProposalToLocalDB").time();
+      UpdateAspectResult result =
+          ingestAspectToLocalDB(entityUrn, metadataChangeProposal.getAspectName(), ignored -> aspect, auditStamp,
+              systemMetadata);
+      ingestToLocalDBTimer.stop();
+      oldAspect = result.getOldValue();
+      oldSystemMetadata = result.getOldSystemMetadata();
+      newAspect = result.getNewValue();
+      newSystemMetadata = result.getNewSystemMetadata();
+      // Apply retention policies asynchronously if there was an update to existing aspect value
+      if (oldAspect != newAspect && oldAspect != null && retentionService != null) {
+        retentionService.applyRetention(entityUrn, aspectSpec.getName(),
+            Optional.of(new RetentionService.RetentionContext(Optional.of(result.maxVersion))));
+      }
+    }
+
+    if (oldAspect != newAspect || getAlwaysEmitAuditEvent()) {
+      log.debug(String.format("Producing MetadataChangeLog for ingested aspect %s, urn %s",
+          metadataChangeProposal.getAspectName(), entityUrn));
+
+      final MetadataChangeLog metadataChangeLog = new MetadataChangeLog(metadataChangeProposal.data());
+      if (oldAspect != null) {
+        metadataChangeLog.setPreviousAspectValue(GenericAspectUtils.serializeAspect(oldAspect));
+      }
+      if (oldSystemMetadata != null) {
+        metadataChangeLog.setPreviousSystemMetadata(oldSystemMetadata);
+      }
+      if (newAspect != null) {
+        metadataChangeLog.setAspect(GenericAspectUtils.serializeAspect(newAspect));
+      }
+      if (newSystemMetadata != null) {
+        metadataChangeLog.setSystemMetadata(newSystemMetadata);
+      }
+
+      log.debug(String.format("Serialized MCL event: %s", metadataChangeLog));
+      // Since only timeseries aspects are ingested as of now, simply produce mae event for it
+      produceMetadataChangeLog(entityUrn, aspectSpec, metadataChangeLog);
+    } else {
+      log.debug(
+          String.format("Skipped producing MetadataAuditEvent for ingested aspect %s, urn %s. Aspect has not changed.",
+              metadataChangeProposal.getAspectName(), entityUrn));
+    }
+
+    return new IngestProposalResult(entityUrn, oldAspect != newAspect);
+  }
+
   /**
    * Updates a particular version of an aspect & optionally emits a {@link com.linkedin.mxe.MetadataAuditEvent}.
    *
@@ -483,12 +646,12 @@ private Map<String, Set<String>> buildEntityToValidAspects(final EntityRegistry
             entry -> entry.getAspectSpecs().stream().map(AspectSpec::getName).collect(Collectors.toSet())));
   }
 
-  public Boolean getEmitAspectSpecificAuditEvent() {
-    return _emitAspectSpecificAuditEvent;
+  public Boolean getAlwaysEmitAuditEvent() {
+    return _alwaysEmitAuditEvent;
   }
 
-  public void setEmitAspectSpecificAuditEvent(Boolean emitAspectSpecificAuditEvent) {
-    _emitAspectSpecificAuditEvent = emitAspectSpecificAuditEvent;
+  public void setAlwaysEmitAuditEvent(Boolean alwaysEmitAuditEvent) {
+    _alwaysEmitAuditEvent = alwaysEmitAuditEvent;
   }
 
   public EntityRegistry getEntityRegistry() {
@@ -505,8 +668,6 @@ protected Set<String> getEntityAspectNames(final String entityName) {
 
   public abstract void setWritable(boolean canWrite);
 
-  public abstract Urn ingestProposal(MetadataChangeProposal metadataChangeProposal, AuditStamp auditStamp);
-
   public RollbackRunResult rollbackRun(List<AspectRowSummary> aspectRows, String runId) {
     return rollbackWithConditions(aspectRows, Collections.singletonMap("runId", runId));
   }
@@ -517,4 +678,21 @@ public abstract RollbackRunResult rollbackWithConditions(List<AspectRowSummary>
   public abstract RollbackRunResult deleteUrn(Urn urn);
 
   public abstract Boolean exists(Urn urn);
+
+  @Value
+  public static class UpdateAspectResult {
+    Urn urn;
+    RecordTemplate oldValue;
+    RecordTemplate newValue;
+    SystemMetadata oldSystemMetadata;
+    SystemMetadata newSystemMetadata;
+    MetadataAuditOperation operation;
+    long maxVersion;
+  }
+
+  @Value
+  public static class IngestProposalResult {
+    Urn urn;
+    boolean didUpdate;
+  }
 }