Add restore command to ZK CLI

mgao0 · mgao0 · commit cd61c1d448ed · 2021-03-22T15:32:37.000-07:00
diff --git a/zookeeper-server/src/main/java/org/apache/zookeeper/ZooKeeperMain.java b/zookeeper-server/src/main/java/org/apache/zookeeper/ZooKeeperMain.java
@@ -57,6 +57,7 @@
 import org.apache.zookeeper.cli.MalformedCommandException;
 import org.apache.zookeeper.cli.ReconfigCommand;
 import org.apache.zookeeper.cli.RemoveWatchesCommand;
+import org.apache.zookeeper.cli.RestoreCommand;
 import org.apache.zookeeper.cli.SetAclCommand;
 import org.apache.zookeeper.cli.SetCommand;
 import org.apache.zookeeper.cli.SetQuotaCommand;
@@ -123,6 +124,7 @@ public boolean getPrintWatches() {
         new GetAllChildrenNumberCommand().addToMap(commandMapCli);
         new VersionCommand().addToMap(commandMapCli);
         new AddWatchCommand().addToMap(commandMapCli);
+        new RestoreCommand().addToMap(commandMapCli);
 
         // add all to commandMap
         for (Entry<String, CliCommand> entry : commandMapCli.entrySet()) {
diff --git a/zookeeper-server/src/main/java/org/apache/zookeeper/cli/CliParseException.java b/zookeeper-server/src/main/java/org/apache/zookeeper/cli/CliParseException.java
@@ -31,4 +31,7 @@ public CliParseException(String message) {
         super(message);
     }
 
+  public CliParseException(String message, Throwable cause) {
+    super(message, cause);
+  }
 }
diff --git a/zookeeper-server/src/main/java/org/apache/zookeeper/cli/RestoreCommand.java b/zookeeper-server/src/main/java/org/apache/zookeeper/cli/RestoreCommand.java
@@ -0,0 +1,131 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+
+package org.apache.zookeeper.cli;
+
+import org.apache.commons.cli.CommandLine;
+import org.apache.commons.cli.Option;
+import org.apache.commons.cli.Options;
+import org.apache.commons.cli.ParseException;
+import org.apache.commons.cli.Parser;
+import org.apache.commons.cli.PosixParser;
+import org.apache.zookeeper.server.backup.RestoreFromBackupTool;
+
+/**
+ * Restore command for ZkCli.
+ */
+public class RestoreCommand extends CliCommand {
+
+  private RestoreFromBackupTool tool;
+  private CommandLine cl;
+  private static Options options = new Options();
+  private static final String RESTORE_ZXID_STR = "restore_zxid";
+  private static final String RESTORE_TIMESTAMP_STR = "restore_timestamp";
+  private static final String BACKUP_STORE_STR = "backup_store";
+  private static final String DATA_DESTINATION_STR = "data_destination";
+  private static final String LOG_DESTINATION_STR = "log_destination";
+  private static final String TIMETABLE_STORAGE_PATH_STR = "timetable_storage_path";
+  private static final String LOCAL_RESTORE_TEMP_DIR_PATH_STR = "local_restore_temp_dir_path";
+  private static final String DRY_RUN_STR = "dry_run";
+  public static final String RESTORE_ZXID_OPTION = "z";
+  public static final String RESTORE_TIMESTAMP_OPTION = "t";
+  public static final String BACKUP_STORE_OPTION = "b";
+  public static final String DATA_DESTINATION_OPTION = "d";
+  public static final String LOG_DESTINATION_OPTION = "l";
+  public static final String TIMETABLE_STORAGE_PATH_OPTION = "s";
+  public static final String LOCAL_RESTORE_TEMP_DIR_PATH_OPTION = "r";
+  public static final String DRY_RUN_OPTION = "n";
+  private static final String RESTORE_ZXID_CMD = "-" + RESTORE_ZXID_OPTION + " " + RESTORE_ZXID_STR;
+  private static final String RESTORE_TIMESTAMP_CMD =
+      "-" + RESTORE_TIMESTAMP_OPTION + " " + RESTORE_TIMESTAMP_STR;
+  private static final String BACKUP_STORE_CMD = "-" + BACKUP_STORE_OPTION + " " + BACKUP_STORE_STR;
+  private static final String DATA_DESTINATION_CMD =
+      "-" + DATA_DESTINATION_OPTION + " " + DATA_DESTINATION_STR;
+  private static final String LOG_DESTINATION_CMD =
+      "-" + LOG_DESTINATION_OPTION + " " + LOG_DESTINATION_STR;
+  private static final String TIMETABLE_STORAGE_PATH_CMD =
+      "-" + TIMETABLE_STORAGE_PATH_OPTION + " " + TIMETABLE_STORAGE_PATH_STR;
+  private static final String LOCAL_RESTORE_TEMP_DIR_PATH_CMD =
+      "-" + LOCAL_RESTORE_TEMP_DIR_PATH_OPTION + " " + LOCAL_RESTORE_TEMP_DIR_PATH_STR;
+  private static final String DRY_RUN_CMD = "-" + DRY_RUN_OPTION;
+  private static final String RESTORE_CMD_STR = "restore";
+  private static final String OPTION_STR =
+      "[" + RESTORE_ZXID_CMD + "]/[" + RESTORE_TIMESTAMP_CMD + "] [" + BACKUP_STORE_CMD + "] ["
+          + DATA_DESTINATION_CMD + "] [" + LOG_DESTINATION_CMD + "] [" + TIMETABLE_STORAGE_PATH_CMD
+          + "](needed if restore to a timestamp) [" + LOCAL_RESTORE_TEMP_DIR_PATH_CMD
+          + "](optional) [" + DRY_RUN_CMD + "](optional)";
+
+  static {
+    options.addOption(new Option(RESTORE_ZXID_OPTION, true, RESTORE_ZXID_STR));
+    options.addOption(new Option(RESTORE_TIMESTAMP_OPTION, true, RESTORE_TIMESTAMP_STR));
+    options.addOption(new Option(BACKUP_STORE_OPTION, true, BACKUP_STORE_STR));
+    options.addOption(new Option(DATA_DESTINATION_OPTION, true, DATA_DESTINATION_STR));
+    options.addOption(new Option(LOG_DESTINATION_OPTION, true, LOG_DESTINATION_STR));
+    options.addOption(new Option(TIMETABLE_STORAGE_PATH_OPTION, true, TIMETABLE_STORAGE_PATH_STR));
+    options.addOption(
+        new Option(LOCAL_RESTORE_TEMP_DIR_PATH_OPTION, true, LOCAL_RESTORE_TEMP_DIR_PATH_STR));
+    options.addOption(new Option(DRY_RUN_OPTION, false, DRY_RUN_STR));
+  }
+
+  public RestoreCommand() {
+    super(RESTORE_CMD_STR, OPTION_STR);
+    tool = new RestoreFromBackupTool();
+  }
+
+  @Override
+  public String getUsageStr() {
+    return "Usage: RestoreFromBackupTool  " + RESTORE_CMD_STR + " " + OPTION_STR + "\n    "
+        + RESTORE_ZXID_CMD
+        + ": the point to restore to, either the string 'latest' or a zxid in hex format. Choose one between this option or "
+        + RESTORE_TIMESTAMP_CMD + ", if both are specified, this option will be prioritized\n    "
+        + RESTORE_TIMESTAMP_CMD
+        + ": the point to restore to, a timestamp in long format. Choose one between this option or "
+        + RESTORE_ZXID_CMD + ".\n    " + BACKUP_STORE_CMD
+        + ": the connection information for the backup store\n           For GPFS the format is: gpfs:<config_path>:<backup_path>:<namespace>\n    "
+        + DATA_DESTINATION_CMD + ": local destination path for restored snapshots\n    "
+        + LOG_DESTINATION_CMD + ": local destination path for restored txlogs\n    "
+        + TIMETABLE_STORAGE_PATH_CMD
+        + ": Needed if restore to a timestamp. Backup storage path for timetable files, for GPFS the format is: gpfs:<config_path>:<backup_path>:<namespace>, if not set, default to be same as backup storage path\n    "
+        + LOCAL_RESTORE_TEMP_DIR_PATH_CMD
+        + ": Optional, local path for creating a temporary intermediate directory for restoration, the directory will be deleted after restoration is done\n    "
+        + DRY_RUN_CMD + " " + DRY_RUN_STR
+        + ": Optional, no files will be actually copied in a dry run";
+  }
+
+  @Override
+  public CliCommand parse(String[] cmdArgs) throws CliParseException {
+    Parser parser = new PosixParser();
+    try {
+      cl = parser.parse(options, cmdArgs);
+    } catch (ParseException ex) {
+      throw new CliParseException(getUsageStr(), ex);
+    }
+    if ((!cl.hasOption(RESTORE_ZXID_OPTION) && !cl.hasOption(RESTORE_TIMESTAMP_OPTION)) || !cl
+        .hasOption(BACKUP_STORE_OPTION) || !cl.hasOption(DATA_DESTINATION_OPTION) || !cl
+        .hasOption(LOG_DESTINATION_OPTION)) {
+      throw new CliParseException("Missing required argument(s).\n" + getUsageStr());
+    }
+    return this;
+  }
+
+  @Override
+  public boolean exec() throws CliException {
+    return tool.runWithRetries(cl);
+  }
+}
diff --git a/zookeeper-server/src/main/java/org/apache/zookeeper/server/backup/BackupConfig.java b/zookeeper-server/src/main/java/org/apache/zookeeper/server/backup/BackupConfig.java
@@ -68,8 +68,8 @@ public class BackupConfig {
   private final String backupStoragePath;
   /*
    * The custom namespace string under which the backup files will be stored in.
-   * E.g.) <backupStoragePath>/<namespace>/snapshot/snapshot-123456
-   *       <backupStoragePath>/<namespace>/translog/translog-123456
+   * E.g.) <backupStoragePath>/<namespace>/snapshot-123456
+   *       <backupStoragePath>/<namespace>/translog-123456
    */
   private final String namespace;
 
@@ -93,7 +93,7 @@ public BackupConfig(Builder builder) {
     this.storageProviderClassName = builder.storageProviderClassName.get();
     this.storageConfig = builder.storageConfig.orElse(null);
     this.backupStoragePath = builder.backupStoragePath.orElse("");
-    this.namespace = builder.namespace.orElse("");
+    this.namespace = builder.namespace.orElse("UNKNOWN");
     this.timetableEnabled = builder.timetableEnabled.orElse(false);
     // If timetable storage path is not given, use the backup storage path
     this.timetableStoragePath = builder.timetableStoragePath.orElse(backupStoragePath);
@@ -405,4 +405,45 @@ private void validate() throws ConfigException {
       }
     }
   }
+
+  public static class RestorationBackupConfigBuilder extends Builder {
+    private Optional<String> storageProviderClassName = Optional.empty();
+    private Optional<File> storageConfig = Optional.empty();
+    private Optional<String> backupStoragePath = Optional.empty();
+    private Optional<String> namespace = Optional.empty();
+
+    public RestorationBackupConfigBuilder setStorageProviderClassName(
+        String storageProviderClassName) {
+      this.storageProviderClassName = Optional.of(storageProviderClassName);
+      return this;
+    }
+
+    public RestorationBackupConfigBuilder setStorageConfig(File storageConfig) {
+      this.storageConfig = Optional.of(storageConfig);
+      return this;
+    }
+
+    public RestorationBackupConfigBuilder setBackupStoragePath(String backupStoragePath) {
+      this.backupStoragePath = Optional.of(backupStoragePath);
+      return this;
+    }
+
+    public RestorationBackupConfigBuilder setNamespace(String namespace) {
+      this.namespace = Optional.of(namespace);
+      return this;
+    }
+
+    public Optional<BackupConfig> build() throws ConfigException {
+      if (!storageProviderClassName.isPresent() || storageProviderClassName.get().isEmpty()) {
+        throw new ConfigException("Please specify a valid storage provider class name.");
+      }
+      if (!backupStoragePath.isPresent()) {
+        throw new ConfigException("Please specify a valid backup storage path..");
+      }
+      if (!namespace.isPresent()) {
+        throw new ConfigException("Please specify a valid namespace.");
+      }
+      return Optional.of(new BackupConfig(this));
+    }
+  }
 }
diff --git a/zookeeper-server/src/main/java/org/apache/zookeeper/server/backup/BackupManager.java b/zookeeper-server/src/main/java/org/apache/zookeeper/server/backup/BackupManager.java
@@ -503,7 +503,7 @@ public BackupManager(File snapDir, File dataLogDir, long serverId, BackupConfig
     this.serverId = serverId;
     this.namespace = backupConfig.getNamespace() == null ? "UNKNOWN" : backupConfig.getNamespace();
     try {
-      backupStorage = createStorageProviderImpl(backupConfig);
+      backupStorage = BackupUtil.createStorageProviderImpl(backupConfig);
     } catch (ReflectiveOperationException e) {
       throw new BackupException(e.getMessage(), e);
     }
@@ -597,7 +597,7 @@ public synchronized void initialize() throws IOException {
       // This is because we want the timetable backup to be stored in a different storage path
       BackupStorageProvider timetableBackupStorage;
       try {
-        timetableBackupStorage = createStorageProviderImpl(
+        timetableBackupStorage = BackupUtil.createStorageProviderImpl(
             backupConfig.getBuilder().setBackupStoragePath(backupConfig.getTimetableStoragePath())
                 .build().get());
         LOG.info(
@@ -612,25 +612,4 @@ public synchronized void initialize() throws IOException {
       timetableBackup.initialize();
     }
   }
-
-  /**
-   * Instantiates the storage provider implementation by reflection. This allows the user to
-   * choose which BackupStorageProvider implementation to use by specifying the fully-qualified
-   * class name in BackupConfig (read from Properties).
-   * @param backupConfig
-   * @return
-   * @throws ClassNotFoundException
-   * @throws NoSuchMethodException
-   * @throws IllegalAccessException
-   * @throws InvocationTargetException
-   * @throws InstantiationException
-   */
-  private BackupStorageProvider createStorageProviderImpl(BackupConfig backupConfig)
-      throws ClassNotFoundException, NoSuchMethodException, IllegalAccessException,
-             InvocationTargetException, InstantiationException {
-    Class<?> clazz = Class.forName(backupConfig.getStorageProviderClassName());
-    Constructor<?> constructor = clazz.getConstructor(BackupConfig.class);
-    Object storageProvider = constructor.newInstance(backupConfig);
-    return (BackupStorageProvider) storageProvider;
-  }
 }
diff --git a/zookeeper-server/src/main/java/org/apache/zookeeper/server/backup/BackupUtil.java b/zookeeper-server/src/main/java/org/apache/zookeeper/server/backup/BackupUtil.java
@@ -21,6 +21,8 @@
 import java.io.File;
 import java.io.IOException;
 import java.io.Serializable;
+import java.lang.reflect.Constructor;
+import java.lang.reflect.InvocationTargetException;
 import java.util.*;
 
 import com.google.common.base.Function;
@@ -356,6 +358,27 @@ public static List<Range<Long>> getRanges(
             zxidRangeExtractor));
   }
 
+  /**
+   * Instantiates the storage provider implementation by reflection. This allows the user to
+   * choose which BackupStorageProvider implementation to use by specifying the fully-qualified
+   * class name in BackupConfig (read from Properties).
+   * @param backupConfig
+   * @return
+   * @throws ClassNotFoundException
+   * @throws NoSuchMethodException
+   * @throws IllegalAccessException
+   * @throws InvocationTargetException
+   * @throws InstantiationException
+   */
+  public static BackupStorageProvider createStorageProviderImpl(BackupConfig backupConfig)
+      throws ClassNotFoundException, NoSuchMethodException, IllegalAccessException,
+             InvocationTargetException, InstantiationException {
+    Class<?> clazz = Class.forName(backupConfig.getStorageProviderClassName());
+    Constructor<?> constructor = clazz.getConstructor(BackupConfig.class);
+    Object storageProvider = constructor.newInstance(backupConfig);
+    return (BackupStorageProvider) storageProvider;
+  }
+
   // Utility class
   private BackupUtil() {}
 }
diff --git a/zookeeper-server/src/main/java/org/apache/zookeeper/server/backup/RestoreFromBackupTool.java b/zookeeper-server/src/main/java/org/apache/zookeeper/server/backup/RestoreFromBackupTool.java
diff --git a/zookeeper-server/src/test/java/org/apache/zookeeper/server/backup/RestorationToolTest.java b/zookeeper-server/src/test/java/org/apache/zookeeper/server/backup/RestorationToolTest.java

Original file line number	Diff line number	Diff line change
`@@ -31,4 +31,7 @@ public CliParseException(String message) {`
`31`	`31`	`super(message);`
`32`	`32`	`}`
`33`	`33`
	`34`	`+ public CliParseException(String message, Throwable cause) {`
	`35`	`+ super(message, cause);`
	`36`	`+ }`
`34`	`37`	`}`