diff --git a/tools/ldb_cmd.cc b/tools/ldb_cmd.cc index d0878f894..b56eb4c34 100644 --- a/tools/ldb_cmd.cc +++ b/tools/ldb_cmd.cc @@ -223,6 +223,9 @@ LDBCommand* LDBCommand::SelectCommand(const ParsedParams& parsed_params) { } else if (parsed_params.cmd == BackupCommand::Name()) { return new BackupCommand(parsed_params.cmd_params, parsed_params.option_map, parsed_params.flags); + } else if (parsed_params.cmd == RestoreCommand::Name()) { + return new RestoreCommand(parsed_params.cmd_params, + parsed_params.option_map, parsed_params.flags); } return nullptr; } @@ -2545,10 +2548,7 @@ BackupCommand::BackupCommand(const std::vector& params, thread_num_ = std::stoi(itr->second); } itr = options.find(ARG_BACKUP_ENV); - if (itr == options.end()) { - exec_state_ = LDBCommandExecuteResult::Failed("--" + ARG_BACKUP_ENV + - ": missing backup env"); - } else { + if (itr != options.end()) { test_cluster_ = itr->second; } itr = options.find(ARG_BACKUP_DIR); @@ -2600,6 +2600,72 @@ void BackupCommand::DoCommand() { // ---------------------------------------------------------------------------- +const std::string RestoreCommand::ARG_BACKUP_ENV_URI = "backup_env_uri"; +const std::string RestoreCommand::ARG_BACKUP_DIR = "backup_dir"; +const std::string RestoreCommand::ARG_NUM_THREADS = "num_threads"; + +RestoreCommand::RestoreCommand( + const std::vector& params, + const std::map& options, + const std::vector& flags) + : LDBCommand(options, flags, false /* is_read_only */, + BuildCmdLineOptions( + {ARG_BACKUP_ENV_URI, ARG_BACKUP_DIR, ARG_NUM_THREADS})), + num_threads_(1) { + auto itr = options.find(ARG_NUM_THREADS); + if (itr != options.end()) { + num_threads_ = std::stoi(itr->second); + } + itr = options.find(ARG_BACKUP_ENV_URI); + if (itr != options.end()) { + backup_env_uri_ = itr->second; + } + itr = options.find(ARG_BACKUP_DIR); + if (itr == options.end()) { + exec_state_ = LDBCommandExecuteResult::Failed("--" + ARG_BACKUP_DIR + + ": missing backup directory"); + } else { + backup_dir_ = itr->second; + } +} + +void RestoreCommand::Help(std::string& ret) { + ret.append(" "); + ret.append(RestoreCommand::Name()); + ret.append(" [--" + ARG_BACKUP_ENV_URI + "] "); + ret.append(" [--" + ARG_BACKUP_DIR + "] "); + ret.append(" [--" + ARG_NUM_THREADS + "] "); + ret.append("\n"); +} + +void RestoreCommand::DoCommand() { + std::unique_ptr custom_env_guard; + Env* custom_env = NewEnvFromUri(backup_env_uri_, &custom_env_guard); + std::unique_ptr restore_engine; + Status status; + { + BackupableDBOptions opts(backup_dir_, custom_env); + opts.max_background_operations = num_threads_; + BackupEngineReadOnly* raw_restore_engine_ptr; + status = BackupEngineReadOnly::Open(Env::Default(), opts, + &raw_restore_engine_ptr); + if (status.ok()) { + restore_engine.reset(raw_restore_engine_ptr); + } + } + if (status.ok()) { + printf("open restore engine OK\n"); + status = restore_engine->RestoreDBFromLatestBackup(db_path_, db_path_); + } + if (status.ok()) { + printf("restore from backup OK\n"); + } else { + exec_state_ = LDBCommandExecuteResult::Failed(status.ToString()); + } +} + +// ---------------------------------------------------------------------------- + namespace { void DumpSstFile(std::string filename, bool output_hex, bool show_properties) { diff --git a/tools/ldb_cmd_impl.h b/tools/ldb_cmd_impl.h index 6de9534b9..ceae95a76 100644 --- a/tools/ldb_cmd_impl.h +++ b/tools/ldb_cmd_impl.h @@ -468,4 +468,26 @@ class BackupCommand : public LDBCommand { static const std::string ARG_THREAD; }; +class RestoreCommand : public LDBCommand { + public: + static std::string Name() { return "restore"; } + + RestoreCommand(const std::vector& params, + const std::map& options, + const std::vector& flags); + + virtual void DoCommand() override; + virtual bool NoDBOpen() override { return true; } + + static void Help(std::string& ret); + + private: + std::string backup_env_uri_; + std::string backup_dir_; + int num_threads_; + + static const std::string ARG_BACKUP_DIR; + static const std::string ARG_BACKUP_ENV_URI; + static const std::string ARG_NUM_THREADS; +}; } // namespace rocksdb diff --git a/tools/ldb_tool.cc b/tools/ldb_tool.cc index 2f9219c13..7f4f5a9fa 100644 --- a/tools/ldb_tool.cc +++ b/tools/ldb_tool.cc @@ -80,6 +80,7 @@ void LDBCommandRunner::PrintHelp(const char* exec_name) { InternalDumpCommand::Help(ret); RepairCommand::Help(ret); BackupCommand::Help(ret); + RestoreCommand::Help(ret); fprintf(stderr, "%s\n", ret.c_str()); }