Add cleanupData command to nrt_utils (#733)

Yelp · Sep 30, 2024 · 1e75d95 · 1e75d95
1 parent 680bd73
commit 1e75d95
Show file tree

Hide file tree

Showing 6 changed files with 959 additions and 3 deletions.
diff --git a/build.gradle b/build.gradle
@@ -169,11 +169,13 @@ test {
     environment(Map.of('CUSTOM_HOST', 'my_custom_host', 'VAR1', 'v1', 'VAR2', 'v2', 'VAR3', 'v3'))
     if (project.hasProperty('longRunningTestsOnly')) {
         include '**/IncrementalDataCleanupCommandTest.class'
+        include '**/CleanupDataCommandTest.class'
     } else {
         if (!project.hasProperty('includePerfTests')) {
             exclude '**/YelpReviewsTest.class'
             exclude '**/MergeBehaviorTests.class'
             exclude '**/IncrementalDataCleanupCommandTest.class'
+            exclude '**/CleanupDataCommandTest.class'
             filter {
                 excludeTestsMatching '*.NodeNameResolverAndLoadBalancingTests.testSimpleLoadBalancingAsync'
             }

diff --git a/src/main/java/com/yelp/nrtsearch/server/remote/s3/S3Backend.java b/src/main/java/com/yelp/nrtsearch/server/remote/s3/S3Backend.java
@@ -68,7 +68,7 @@ public class S3Backend implements RemoteBackend {
   static final String POINT_STATE = "point_state";
   static final String DATA = "data";
   static final String WARMING = "warming";
-  static final String CURRENT_VERSION = "_current";
+  public static final String CURRENT_VERSION = "_current";
 
   private static final Logger logger = LoggerFactory.getLogger(S3Backend.class);
   private static final String ZIP_EXTENSION = ".zip";

diff --git a/src/main/java/com/yelp/nrtsearch/tools/nrt_utils/NrtUtilsCommand.java b/src/main/java/com/yelp/nrtsearch/tools/nrt_utils/NrtUtilsCommand.java
@@ -19,6 +19,7 @@
 import com.yelp.nrtsearch.tools.nrt_utils.backup.ListSnapshotsCommand;
 import com.yelp.nrtsearch.tools.nrt_utils.backup.RestoreCommand;
 import com.yelp.nrtsearch.tools.nrt_utils.backup.SnapshotCommand;
+import com.yelp.nrtsearch.tools.nrt_utils.cleanup.CleanupDataCommand;
 import com.yelp.nrtsearch.tools.nrt_utils.legacy.incremental.DeleteIncrementalSnapshotsCommand;
 import com.yelp.nrtsearch.tools.nrt_utils.legacy.incremental.IncrementalDataCleanupCommand;
 import com.yelp.nrtsearch.tools.nrt_utils.legacy.incremental.ListIncrementalSnapshotsCommand;
@@ -36,6 +37,7 @@
     name = "nrt_utils",
     synopsisSubcommandLabel = "COMMAND",
     subcommands = {
+      CleanupDataCommand.class,
       CleanupSnapshotsCommand.class,
       DeleteIncrementalSnapshotsCommand.class,
       GetRemoteStateCommand.class,

diff --git a/src/main/java/com/yelp/nrtsearch/tools/nrt_utils/backup/BackupCommandUtils.java b/src/main/java/com/yelp/nrtsearch/tools/nrt_utils/backup/BackupCommandUtils.java
@@ -92,7 +92,7 @@ public static String getSnapshotIndexMetadataPrefix(String snapshotRoot, String
    * @param bucketName s3 bucket
    * @param keys keys to delete
    */
-  static void deleteObjects(AmazonS3 s3Client, String bucketName, List<String> keys) {
+  public static void deleteObjects(AmazonS3 s3Client, String bucketName, List<String> keys) {
     System.out.println("Batch deleting objects, size: " + keys.size());
     DeleteObjectsRequest multiObjectDeleteRequest =
         new DeleteObjectsRequest(bucketName).withKeys(keys.toArray(new String[0])).withQuiet(true);
@@ -107,7 +107,7 @@ static void deleteObjects(AmazonS3 s3Client, String bucketName, List<String> key
    * @param interval interval string
    * @return interval in ms
    */
-  static long getTimeIntervalMs(String interval) {
+  public static long getTimeIntervalMs(String interval) {
     String trimmed = interval.trim();
     if (trimmed.length() < 2) {
       throw new IllegalArgumentException("Invalid time interval: " + trimmed);