feat(retention): change default policy, add API to apply retention

.github/pr-labeler-config.yml

@@ -9,6 +9,7 @@ product:
 - 'datahub-web-react/**/*'
 - 'datahub-frontend/**/*'
 - 'datahub-graphql-core/**/*'
+- 'metadata-io/**/*'
 
 docs:
 - 'docs/**/*'

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

@@ -41,6 +41,8 @@
 import com.linkedin.metadata.entity.ebean.EbeanAspectV2;
 import com.linkedin.metadata.entity.restoreindices.RestoreIndicesArgs;
 import com.linkedin.metadata.entity.restoreindices.RestoreIndicesResult;
+import com.linkedin.metadata.entity.retention.BulkApplyRetentionArgs;
+import com.linkedin.metadata.entity.retention.BulkApplyRetentionResult;
 import com.linkedin.metadata.entity.validation.EntityRegistryUrnValidator;
 import com.linkedin.metadata.entity.validation.RecordTemplateValidator;
 import com.linkedin.metadata.entity.validation.ValidationUtils;
@@ -986,6 +988,27 @@ private UpdateAspectResult patchAspect(final Patch patch, final SystemMetadata s
     return result;
   }
 
+  public String batchApplyRetention(Integer start, Integer count, Integer attemptWithVersion, String aspectName,
+                                    String urn) {
+    BulkApplyRetentionArgs args = new BulkApplyRetentionArgs();
+    if (start == null) {
+      start = 0;
+    }
+    args.start = start;
+    if (count == null) {
+      count = 100;
+    }
+    args.count = count;
+    if (attemptWithVersion == null) {
+      attemptWithVersion = 21;
+    }
+    args.attemptWithVersion = attemptWithVersion;
+    args.aspectName = aspectName;
+    args.urn = urn;
+    BulkApplyRetentionResult result = _retentionService.batchApplyRetentionEntities(args);
+    return result.toString();
+  }
+
   private boolean emitChangeLog(@Nullable RecordTemplate oldAspect, @Nullable SystemMetadata oldSystemMetadata,
       RecordTemplate newAspect, SystemMetadata newSystemMetadata,
       MetadataChangeProposal mcp, Urn entityUrn,

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

@@ -7,6 +7,8 @@
 import com.linkedin.data.template.RecordTemplate;
 import com.linkedin.events.metadata.ChangeType;
 import com.linkedin.metadata.Constants;
+import com.linkedin.metadata.entity.retention.BulkApplyRetentionArgs;
+import com.linkedin.metadata.entity.retention.BulkApplyRetentionResult;
 import com.linkedin.metadata.key.DataHubRetentionKey;
 import com.linkedin.metadata.utils.EntityKeyUtils;
 import com.linkedin.metadata.utils.GenericRecordUtils;
@@ -184,6 +186,12 @@ public abstract void applyRetention(@Nonnull Urn urn, @Nonnull String aspectName
    */
   public abstract void batchApplyRetention(@Nullable String entityName, @Nullable String aspectName);
 
+  /**
+   * Batch apply retention to all records within the start, end count
+   */
+  public abstract BulkApplyRetentionResult batchApplyRetentionEntities(@Nonnull BulkApplyRetentionArgs args);
+
+
   @Value
   public static class RetentionContext {
     Optional<Long> maxVersion;

metadata-io/src/main/java/com/linkedin/metadata/entity/cassandra/CassandraRetentionService.java

@@ -14,6 +14,8 @@
 import com.linkedin.metadata.entity.RetentionService;
 import com.linkedin.metadata.entity.EntityAspectIdentifier;
 import com.linkedin.metadata.entity.EntityAspect;
+import com.linkedin.metadata.entity.retention.BulkApplyRetentionArgs;
+import com.linkedin.metadata.entity.retention.BulkApplyRetentionResult;
 import com.linkedin.retention.DataHubRetentionConfig;
 import com.linkedin.retention.Retention;
 import com.linkedin.retention.TimeBasedRetention;
@@ -112,6 +114,12 @@ public void batchApplyRetention(@Nullable String entityName, @Nullable String as
     log.info("Finished applying retention to all records");
   }
 
+  @Override
+  public BulkApplyRetentionResult batchApplyRetentionEntities(@Nonnull BulkApplyRetentionArgs args) {
+    log.error("batchApplyRetentionEntities not implemented for cassandra");
+    return null;
+  }
+
   private void applyVersionBasedRetention(
       @Nonnull final Urn urn,
       @Nonnull final String aspectName,

metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/EbeanRetentionService.java

@@ -4,6 +4,8 @@
 import com.datahub.util.RecordUtils;
 import com.linkedin.metadata.entity.EntityService;
 import com.linkedin.metadata.entity.RetentionService;
+import com.linkedin.metadata.entity.retention.BulkApplyRetentionArgs;
+import com.linkedin.metadata.entity.retention.BulkApplyRetentionResult;
 import com.linkedin.retention.DataHubRetentionConfig;
 import com.linkedin.retention.Retention;
 import com.linkedin.retention.TimeBasedRetention;
@@ -12,6 +14,7 @@
 import io.ebean.Expression;
 import io.ebean.ExpressionList;
 import io.ebean.PagedList;
+import io.ebean.Query;
 import io.ebean.Transaction;
 import io.ebeaninternal.server.expression.Op;
 import io.ebeaninternal.server.expression.SimpleExpression;
@@ -111,52 +114,128 @@ private Expression getTimeBasedRetentionQuery(@Nonnull final TimeBasedRetention
         new Timestamp(_clock.millis() - retention.getMaxAgeInSeconds() * 1000));
   }
 
+  private void applyRetention(
+          PagedList<EbeanAspectV2> rows,
+          Map<String, DataHubRetentionConfig> retentionPolicyMap,
+          BulkApplyRetentionResult applyRetentionResult
+  ) {
+    try (Transaction transaction = _server.beginTransaction()) {
+      transaction.setBatchMode(true);
+      transaction.setBatchSize(_batchSize);
+      for (EbeanAspectV2 row : rows.getList()) {
+        // Only run for cases where there's multiple versions of the aspect
+        if (row.getVersion() == 0) {
+          continue;
+        }
+        // 1. Extract an Entity type from the entity Urn
+        Urn urn;
+        try {
+          urn = Urn.createFromString(row.getUrn());
+        } catch (Exception e) {
+          log.error("Failed to serialize urn {}", row.getUrn(), e);
+          continue;
+        }
+        final String aspectNameFromRecord = row.getAspect();
+        log.debug("Handling urn {} aspect {}", row.getUrn(), row.getAspect());
+        // Get the retention policies to apply from the local retention policy map
+        Optional<Retention> retentionPolicy = getRetentionKeys(urn.getEntityType(), aspectNameFromRecord).stream()
+                .map(key -> retentionPolicyMap.get(key.toString()))
+                .filter(Objects::nonNull)
+                .findFirst()
+                .map(DataHubRetentionConfig::getRetention);
+        retentionPolicy.ifPresent(retention -> applyRetention(urn, aspectNameFromRecord, retention,
+                Optional.of(new RetentionContext(Optional.of(row.getVersion())))));
+        if (applyRetentionResult != null) {
+          applyRetentionResult.rowsHandled += 1;
+        }
+      }
+      transaction.commit();
+    }
+  }
+
   @Override
   @WithSpan
   public void batchApplyRetention(@Nullable String entityName, @Nullable String aspectName) {
     log.debug("Applying retention to all records");
-    int numCandidates = queryCandidates(entityName, aspectName).findCount();
+    int numCandidates = queryCandidates(null, entityName, aspectName).findCount();
     log.info("Found {} urn, aspect pair with more than 1 version", numCandidates);
     Map<String, DataHubRetentionConfig> retentionPolicyMap = getAllRetentionPolicies();
 
     int start = 0;
     while (start < numCandidates) {
       log.info("Applying retention to pairs {} through {}", start, start + _batchSize);
       PagedList<EbeanAspectV2> rows = getPagedAspects(entityName, aspectName, start, _batchSize);
+      applyRetention(rows, retentionPolicyMap, null);
+      start += _batchSize;
+    }
 
-      try (Transaction transaction = _server.beginTransaction()) {
-        transaction.setBatchMode(true);
-        transaction.setBatchSize(_batchSize);
-        for (EbeanAspectV2 row : rows.getList()) {
-          // Only run for cases where there's multiple versions of the aspect
-          if (row.getVersion() == 0) {
-            continue;
-          }
-          // 1. Extract an Entity type from the entity Urn
-          Urn urn;
-          try {
-            urn = Urn.createFromString(row.getUrn());
-          } catch (Exception e) {
-            log.error("Failed to serialize urn {}", row.getUrn(), e);
-            continue;
-          }
-          final String aspectNameFromRecord = row.getAspect();
-          // Get the retention policies to apply from the local retention policy map
-          Optional<Retention> retentionPolicy = getRetentionKeys(urn.getEntityType(), aspectNameFromRecord).stream()
-              .map(key -> retentionPolicyMap.get(key.toString()))
-              .filter(Objects::nonNull)
-              .findFirst()
-              .map(DataHubRetentionConfig::getRetention);
-          retentionPolicy.ifPresent(retention -> applyRetention(urn, aspectNameFromRecord, retention,
-              Optional.of(new RetentionContext(Optional.of(row.getVersion())))));
-        }
-        transaction.commit();
+    log.info("Finished applying retention to all records");
+  }
+
+  @Override
+  public BulkApplyRetentionResult batchApplyRetentionEntities(@Nonnull BulkApplyRetentionArgs args) {
+    long startTime = System.currentTimeMillis();
+
+    BulkApplyRetentionResult result = new BulkApplyRetentionResult();
+    result.argStart = args.start;
+    result.argCount = args.count;
+    result.argAttemptWithVersion = args.attemptWithVersion;
+    result.argAspectName = args.aspectName;
+    result.argUrn = args.urn;
+
+    Map<String, DataHubRetentionConfig> retentionPolicyMap = getAllRetentionPolicies();
+    result.timeRetentionPolicyMapMs = System.currentTimeMillis() - startTime;
+    startTime = System.currentTimeMillis();
+
+    //only supports version based retention for batch apply
+    //find urn, aspect pair where distinct versions > 20 to apply retention policy
+    Query<EbeanAspectV2> query = _server.find(EbeanAspectV2.class)
+            .setDistinct(true)
+            .select(String.format(
+                    "%s, %s, count(%s)", EbeanAspectV2.URN_COLUMN, EbeanAspectV2.ASPECT_COLUMN, EbeanAspectV2.VERSION_COLUMN)
+            );
+    ExpressionList<EbeanAspectV2> exp = null;
+    if (args.urn != null || args.aspectName != null) {
+      exp = query.where();
+      if (args.aspectName != null) {
+        exp = exp.eq(EbeanAspectV2.ASPECT_COLUMN, args.aspectName);
       }
+      if (args.urn != null) {
+        exp = exp.eq(EbeanAspectV2.URN_COLUMN, args.urn);
+      }
+    }
+    if (exp == null) {
+      exp = query.having();
+    } else {
+      exp = exp.having();
+    }
 
-      start += _batchSize;
+    PagedList<EbeanAspectV2> rows = exp
+              .gt(String.format("count(%s)", EbeanAspectV2.VERSION_COLUMN), args.attemptWithVersion)
+            .setFirstRow(args.start)
+            .setMaxRows(args.count)
+            .findPagedList();
+    result.timeRowMs = System.currentTimeMillis() - startTime;
+
+    for (EbeanAspectV2 row : rows.getList()) {
+      startTime = System.currentTimeMillis();
+      log.debug("For {},{} version count is {}", row.getUrn(), row.getAspect(), row.getVersion());
+      try {
+        Urn.createFromString(row.getUrn());
+      } catch (Exception e) {
+        log.error("Failed to serialize urn {}", row.getUrn(), e);
+        continue;
+      }
+      PagedList<EbeanAspectV2> rowsToChange = queryCandidates(row.getUrn(), null, row.getAspect())
+              .setFirstRow(args.start)
+              .setMaxRows(args.count)
+              .findPagedList();
+
+      applyRetention(rowsToChange, retentionPolicyMap, result);
+      result.timeApplyRetentionMs += System.currentTimeMillis() - startTime;
     }
 
-    log.info("Finished applying retention to all records");
+    return result;
   }
 
   private Map<String, DataHubRetentionConfig> getAllRetentionPolicies() {
@@ -172,12 +251,16 @@ private Map<String, DataHubRetentionConfig> getAllRetentionPolicies() {
             row -> RecordUtils.toRecordTemplate(DataHubRetentionConfig.class, row.getMetadata())));
   }
 
-  private ExpressionList<EbeanAspectV2> queryCandidates(@Nullable String entityName, @Nullable String aspectName) {
+  private ExpressionList<EbeanAspectV2> queryCandidates(@Nullable String urn,
+          @Nullable String entityName, @Nullable String aspectName) {
     ExpressionList<EbeanAspectV2> query = _server.find(EbeanAspectV2.class)
         .setDistinct(true)
         .select(String.format("%s, %s, max(%s)", EbeanAspectV2.URN_COLUMN, EbeanAspectV2.ASPECT_COLUMN,
             EbeanAspectV2.VERSION_COLUMN))
         .where();
+    if (urn != null) {
+      query.eq(EbeanAspectV2.URN_COLUMN, urn);
+    }
     if (entityName != null) {
       query.like(EbeanAspectV2.URN_COLUMN, String.format("urn:li:%s%%", entityName));
     }
@@ -189,7 +272,7 @@ private ExpressionList<EbeanAspectV2> queryCandidates(@Nullable String entityNam
 
   private PagedList<EbeanAspectV2> getPagedAspects(@Nullable String entityName, @Nullable String aspectName,
       final int start, final int pageSize) {
-    return queryCandidates(entityName, aspectName).orderBy(
+    return queryCandidates(null, entityName, aspectName).orderBy(
         EbeanAspectV2.URN_COLUMN + ", " + EbeanAspectV2.ASPECT_COLUMN)
         .setFirstRow(start)
         .setMaxRows(pageSize)

metadata-io/src/main/java/com/linkedin/metadata/entity/retention/BulkApplyRetentionArgs.java

@@ -0,0 +1,12 @@
+package com.linkedin.metadata.entity.retention;
+
+import lombok.Data;
+
+@Data
+public class BulkApplyRetentionArgs {
+    public Integer start;
+    public Integer count;
+    public Integer attemptWithVersion;
+    public String aspectName;
+    public String urn;
+}

metadata-io/src/main/java/com/linkedin/metadata/entity/retention/BulkApplyRetentionResult.java

@@ -0,0 +1,16 @@
+package com.linkedin.metadata.entity.retention;
+
+import lombok.Data;
+
+@Data
+public class BulkApplyRetentionResult {
+    public long argStart;
+    public long argCount;
+    public long argAttemptWithVersion;
+    public String argUrn;
+    public String argAspectName;
+    public long rowsHandled = 0;
+    public long timeRetentionPolicyMapMs;
+    public long timeRowMs;
+    public long timeApplyRetentionMs = 0;
+}

metadata-service/restli-api/src/main/idl/com.linkedin.entity.entities.restspec.json

@@ -27,6 +27,30 @@
       } ]
     } ],
     "actions" : [ {
+      "name" : "applyRetention",
+      "parameters" : [ {
+        "name" : "start",
+        "type" : "int",
+        "optional" : true
+      }, {
+        "name" : "count",
+        "type" : "int",
+        "optional" : true
+      }, {
+        "name" : "attemptWithVersion",
+        "type" : "int",
+        "optional" : true
+      }, {
+        "name" : "aspectName",
+        "type" : "string",
+        "optional" : true
+      }, {
+        "name" : "urn",
+        "type" : "string",
+        "optional" : true
+      } ],
+      "returns" : "string"
+    }, {
       "name" : "autocomplete",
       "parameters" : [ {
         "name" : "entity",