[PLAT-14460] Implementing restore from s3 backup

Summary: Diff implements support for restoring from s3 backup. Based on the storage config and backup directory provided we query for the most recent backup in that folder and restore it to YBA. Afterwards, we perform a shutdown so any migrations and postgres connections are reset.

Test Plan: Restore from s3 backup, ensure success

Reviewers: vkumar, sanketh, dshubin, vbansal

Reviewed By: dshubin

Subscribers: yugaware

Differential Revision: https://phorge.dev.yugabyte.com/D38245

managed/src/main/java/com/yugabyte/yw/commissioner/tasks/CreateYbaBackup.java

@@ -101,43 +101,45 @@ public void runScheduledBackup(
   public void run() {
     log.info("Execution of CreateYbaBackup");
     CreateYbaBackup.Params taskParams = taskParams();
-    if (taskParams.storageConfigUUID != null) {
-      log.debug("Creating platform backup...");
-      ShellResponse response =
-          replicationHelper.runCommand(replicationManager.new CreatePlatformBackupParams());
-
-      if (response.code != 0) {
-        log.error("Backup failed: " + response.message);
-        throw new PlatformServiceException(
-            INTERNAL_SERVER_ERROR, "Backup failed: " + response.message);
-      }
-      Optional<File> backupOpt = replicationHelper.getMostRecentBackup();
-      if (!backupOpt.isPresent()) {
-        throw new PlatformServiceException(INTERNAL_SERVER_ERROR, "could not find backup file");
-      }
-      File backup = backupOpt.get();
-      CustomerConfig customerConfig = CustomerConfig.get(taskParams.storageConfigUUID);
-      if (customerConfig == null) {
-        throw new PlatformServiceException(
-            INTERNAL_SERVER_ERROR,
-            "Could not find customer config with provided storage config UUID.");
-      }
-      CloudUtil cloudUtil = cloudUtilFactory.getCloudUtil(customerConfig.getName());
-      if (!cloudUtil.uploadYbaBackup(customerConfig.getDataObject(), backup, taskParams.dirName)) {
-        throw new PlatformServiceException(
-            INTERNAL_SERVER_ERROR, "Could not upload YBA backup to cloud storage.");
-      }
+    if (taskParams.storageConfigUUID == null) {
+      log.info("No storage config UUID set, skipping creation of YBA backup.");
+      return;
+    }
+    log.debug("Creating platform backup...");
+    ShellResponse response =
+        replicationHelper.runCommand(replicationManager.new CreatePlatformBackupParams());
 
-      if (!cloudUtil.cleanupUploadedBackups(customerConfig.getDataObject(), taskParams.dirName)) {
-        log.warn(
-            "Error cleaning up uploaded backups to cloud storage, please delete manually to avoid"
-                + " incurring unexpected costs.");
-      }
+    if (response.code != 0) {
+      log.error("Backup failed: " + response.message);
+      throw new PlatformServiceException(
+          INTERNAL_SERVER_ERROR, "Backup failed: " + response.message);
+    }
+    Optional<File> backupOpt = replicationHelper.getMostRecentBackup();
+    if (!backupOpt.isPresent()) {
+      throw new PlatformServiceException(INTERNAL_SERVER_ERROR, "could not find backup file");
+    }
+    File backup = backupOpt.get();
+    CustomerConfig customerConfig = CustomerConfig.get(taskParams.storageConfigUUID);
+    if (customerConfig == null) {
+      throw new PlatformServiceException(
+          INTERNAL_SERVER_ERROR,
+          "Could not find customer config with provided storage config UUID during create.");
+    }
+    CloudUtil cloudUtil = cloudUtilFactory.getCloudUtil(customerConfig.getName());
+    if (!cloudUtil.uploadYbaBackup(customerConfig.getDataObject(), backup, taskParams.dirName)) {
+      throw new PlatformServiceException(
+          INTERNAL_SERVER_ERROR, "Could not upload YBA backup to cloud storage.");
+    }
 
-      // Cleanup backups
-      replicationHelper.cleanupCreatedBackups();
-      log.info(backup.getAbsolutePath());
+    if (!cloudUtil.cleanupUploadedBackups(customerConfig.getDataObject(), taskParams.dirName)) {
+      log.warn(
+          "Error cleaning up uploaded backups to cloud storage, please delete manually to avoid"
+              + " incurring unexpected costs.");
     }
+
+    // Cleanup backups
+    replicationHelper.cleanupCreatedBackups();
+    log.info(backup.getAbsolutePath());
     return;
   }
 }

managed/src/main/java/com/yugabyte/yw/commissioner/tasks/RestoreContinuousBackup.java

@@ -10,19 +10,40 @@
 
 package com.yugabyte.yw.commissioner.tasks;
 
+import static play.mvc.Http.Status.INTERNAL_SERVER_ERROR;
+
 import com.google.inject.Inject;
 import com.yugabyte.yw.commissioner.AbstractTaskBase;
 import com.yugabyte.yw.commissioner.BaseTaskDependencies;
+import com.yugabyte.yw.common.CloudUtil;
+import com.yugabyte.yw.common.CloudUtilFactory;
+import com.yugabyte.yw.common.PlatformServiceException;
+import com.yugabyte.yw.common.Util;
+import com.yugabyte.yw.common.ha.PlatformReplicationHelper;
+import com.yugabyte.yw.common.ha.PlatformReplicationManager;
 import com.yugabyte.yw.forms.AbstractTaskParams;
+import com.yugabyte.yw.models.configs.CustomerConfig;
+import java.io.File;
 import java.util.UUID;
 import lombok.extern.slf4j.Slf4j;
 
 @Slf4j
 public class RestoreContinuousBackup extends AbstractTaskBase {
 
+  private final CloudUtilFactory cloudUtilFactory;
+  private final PlatformReplicationHelper replicationHelper;
+  private final PlatformReplicationManager replicationManager;
+
   @Inject
-  protected RestoreContinuousBackup(BaseTaskDependencies baseTaskDependencies) {
+  protected RestoreContinuousBackup(
+      BaseTaskDependencies baseTaskDependencies,
+      PlatformReplicationHelper replicationHelper,
+      PlatformReplicationManager replicationManager,
+      CloudUtilFactory cloudUtilFactory) {
     super(baseTaskDependencies);
+    this.replicationHelper = replicationHelper;
+    this.replicationManager = replicationManager;
+    this.cloudUtilFactory = cloudUtilFactory;
   }
 
   public static class Params extends AbstractTaskParams {
@@ -37,7 +58,34 @@ protected Params taskParams() {
 
   @Override
   public void run() {
-    log.info("Dummy exeuction of RestoreContinuousBackup");
+    log.info("Exeuction of RestoreContinuousBackup");
+    RestoreContinuousBackup.Params taskParams = taskParams();
+    if (taskParams.storageConfigUUID == null) {
+      log.info("No storage config UUID set, skipping restore.");
+    }
+    // Download backup from remote location
+    CustomerConfig customerConfig = CustomerConfig.get(taskParams.storageConfigUUID);
+    if (customerConfig == null) {
+      throw new PlatformServiceException(
+          INTERNAL_SERVER_ERROR,
+          "Could not find customer config with provided storage config UUID during restore.");
+    }
+    CloudUtil cloudUtil = cloudUtilFactory.getCloudUtil(customerConfig.getName());
+    File backup =
+        cloudUtil.downloadYbaBackup(
+            customerConfig.getDataObject(),
+            taskParams.backupDir,
+            replicationHelper.getReplicationDirFor(taskParams.storageConfigUUID.toString()));
+    if (backup == null) {
+      throw new PlatformServiceException(
+          INTERNAL_SERVER_ERROR, "Could not download YBA backup from cloud storage.");
+    }
+    if (!replicationManager.restoreBackup(backup)) {
+      throw new PlatformServiceException(INTERNAL_SERVER_ERROR, "Error restoring backup to YBA");
+    }
+    // Restart YBA to cause changes to take effect
+    // Do we want to manually insert RestoreContinuousBackup task info?
+    Util.shutdownYbaProcess(0);
     return;
   }
 }

managed/src/main/java/com/yugabyte/yw/common/AWSUtil.java

@@ -40,6 +40,7 @@
 import com.amazonaws.services.s3.model.DeleteObjectsRequest;
 import com.amazonaws.services.s3.model.DeleteObjectsRequest.KeyVersion;
 import com.amazonaws.services.s3.model.GetBucketLocationRequest;
+import com.amazonaws.services.s3.model.GetObjectRequest;
 import com.amazonaws.services.s3.model.ListObjectsV2Request;
 import com.amazonaws.services.s3.model.ListObjectsV2Result;
 import com.amazonaws.services.s3.model.PutObjectRequest;
@@ -65,8 +66,11 @@
 import com.yugabyte.yw.models.configs.data.CustomerConfigStorageS3Data.RegionLocations;
 import com.yugabyte.yw.models.helpers.NodeDetails;
 import java.io.File;
+import java.io.FileOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
+import java.nio.file.Files;
+import java.nio.file.Path;
 import java.security.KeyStore;
 import java.security.SecureRandom;
 import java.util.ArrayList;
@@ -78,6 +82,7 @@
 import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -280,6 +285,102 @@ public boolean uploadYbaBackup(CustomerConfigData configData, File backup, Strin
     return true;
   }
 
+  @Override
+  public File downloadYbaBackup(CustomerConfigData configData, String backupDir, Path localDir) {
+    try {
+      maybeDisableCertVerification();
+      CustomerConfigStorageS3Data s3Data = (CustomerConfigStorageS3Data) configData;
+      AmazonS3 client = createS3Client(s3Data);
+      CloudLocationInfo cLInfo =
+          getCloudLocationInfo(YbcBackupUtil.DEFAULT_REGION_STRING, s3Data, "");
+      log.info("Downloading most recent backup in s3://{}/{}", cLInfo.bucket, backupDir);
+
+      // Get all the backups in the specified bucket/directory
+      List<S3ObjectSummary> allBackups = new ArrayList<>();
+      String nextContinuationToken = null;
+      do {
+        ListObjectsV2Result listObjectsResult;
+        ListObjectsV2Request request =
+            new ListObjectsV2Request().withBucketName(cLInfo.bucket).withPrefix(backupDir);
+        if (StringUtils.isNotBlank(nextContinuationToken)) {
+          request.withContinuationToken(nextContinuationToken);
+        }
+        listObjectsResult = client.listObjectsV2(request);
+
+        if (listObjectsResult.getKeyCount() == 0) {
+          break;
+        }
+        nextContinuationToken = null;
+        if (listObjectsResult.isTruncated()) {
+          nextContinuationToken = listObjectsResult.getNextContinuationToken();
+        }
+        allBackups.addAll(listObjectsResult.getObjectSummaries());
+      } while (nextContinuationToken != null);
+
+      // Sort backups by last modified date (most recent first)
+      List<S3ObjectSummary> sortedBackups =
+          allBackups.stream()
+              .sorted((o1, o2) -> o2.getLastModified().compareTo(o1.getLastModified()))
+              .collect(Collectors.toList());
+      if (sortedBackups.isEmpty()) {
+        log.error("Could not find any backups to restore");
+        return null;
+      }
+      // Iterate through until we find a backup
+      S3ObjectSummary backup = null;
+      Matcher match = null;
+      Pattern backupPattern = Pattern.compile("backup_.*\\.tgz");
+      for (S3ObjectSummary bkp : sortedBackups) {
+        match = backupPattern.matcher(bkp.getKey());
+        if (match.find()) {
+          log.info("Downloading backup s3:{}/{}", bkp.getBucketName(), bkp.getKey());
+          backup = bkp;
+          break;
+        }
+      }
+
+      if (backup == null) {
+        log.error("Could not find matching backup, aborting restore.");
+        return null;
+      }
+
+      // Construct full local filepath with same name as remote backup
+      File localFile = localDir.resolve(match.group()).toFile();
+      // Create directory at platformReplication/<storageConfigUUID> if necessary
+      Files.createDirectories(localFile.getParentFile().toPath());
+      GetObjectRequest getRequest = new GetObjectRequest(cLInfo.bucket, backup.getKey());
+      client.getObject(getRequest);
+      try (S3Object s3Object = client.getObject(getRequest);
+          InputStream inputStream = s3Object.getObjectContent();
+          FileOutputStream outputStream = new FileOutputStream(localFile)) {
+
+        // Write the object content to the local file
+        byte[] buffer = new byte[1024];
+        int bytesRead;
+        while ((bytesRead = inputStream.read(buffer)) != -1) {
+          outputStream.write(buffer, 0, bytesRead);
+        }
+
+      } catch (Exception e) {
+        log.error("Error writing S3 object to file: {}", e.getMessage());
+        return null;
+      }
+      if (!localFile.exists() || localFile.length() == 0) {
+        log.error("Local file does not exist or is empty, aborting restore.");
+        return null;
+      }
+      log.info("Downloaded file from S3 to {}", localFile.getCanonicalPath());
+      return localFile;
+    } catch (AmazonS3Exception e) {
+      log.error("Error occurred while getting object in S3: {}", e.getErrorMessage());
+    } catch (Exception e) {
+      log.error("Unexpected exception while getting object in S3: {}", e.getMessage());
+    } finally {
+      maybeEnableCertVerification();
+    }
+    return null;
+  }
+
   @Override
   public boolean cleanupUploadedBackups(CustomerConfigData configData, String backupDir) {
     try {

managed/src/main/java/com/yugabyte/yw/common/CloudUtil.java

@@ -121,6 +121,11 @@ public default boolean cleanupUploadedBackups(CustomerConfigData configData, Str
     return false;
   }
 
+  public default File downloadYbaBackup(
+      CustomerConfigData configData, String backupDir, Path localDir) {
+    return null;
+  }
+
   public default RestorePreflightResponse
       generateYBBackupRestorePreflightResponseWithoutBackupObject(
           AdvancedRestorePreflightParams preflightParams, CustomerConfigData configData) {

managed/src/main/java/com/yugabyte/yw/common/ha/PlatformReplicationHelper.java

@@ -207,7 +207,7 @@ JsonNode getBackupInfoJson(long frequency, boolean isRunning) {
     return Json.newObject().put("frequency_milliseconds", frequency).put("is_running", isRunning);
   }
 
-  Path getReplicationDirFor(String leader) {
+  public Path getReplicationDirFor(String leader) {
     String storagePath = confGetter.getStaticConf().getString(AppConfigHelper.YB_STORAGE_PATH);
     return Paths.get(storagePath, REPLICATION_DIR, leader);
   }