rocksdb/util/ldb_cmd.h
Zheng Shao be9b862d47 ldb: add "ldb load" command
Summary: This command accepts key-value pairs from stdin with the same format of "ldb dump" command.  This allows us to try out different compression algorithms/block sizes easily.

Test Plan: dump, load, dump, verify the data is the same.

Reviewers: dhruba

Reviewed By: dhruba

CC: leveldb

Differential Revision: https://reviews.facebook.net/D7443
2012-12-19 01:45:59 -08:00

331 lines
6.8 KiB
C++

// Copyright (c) 2012 Facebook. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
#ifndef LEVELDB_UTIL_LDB_H_
#define LEVELDB_UTIL_LDB_H_
#include <string>
#include <iostream>
#include <sstream>
#include <stdlib.h>
#include <algorithm>
#include <stdio.h>
#include "leveldb/db.h"
#include "leveldb/env.h"
#include "leveldb/options.h"
#include "leveldb/iterator.h"
#include "leveldb/slice.h"
#include "db/version_set.h"
#include "util/logging.h"
namespace leveldb {
class LDBCommandExecuteResult {
public:
enum State {
EXEC_NOT_STARTED = 0, EXEC_SUCCEED = 1, EXEC_FAILED = 2,
};
LDBCommandExecuteResult() {
state_ = EXEC_NOT_STARTED;
message_ = "";
}
LDBCommandExecuteResult(State state, std::string& msg) {
state_ = state;
message_ = msg;
}
std::string ToString() {
std::string ret;
switch (state_) {
case EXEC_SUCCEED:
break;
case EXEC_FAILED:
ret.append("Failed: ");
break;
case EXEC_NOT_STARTED:
ret.append("Not started: ");
}
if (!message_.empty()) {
ret.append(message_);
}
return ret;
}
void Reset() {
state_ = EXEC_NOT_STARTED;
message_ = "";
}
bool IsSucceed() {
return state_ == EXEC_SUCCEED;
}
bool IsNotStarted() {
return state_ == EXEC_NOT_STARTED;
}
bool IsFailed() {
return state_ == EXEC_FAILED;
}
static LDBCommandExecuteResult SUCCEED(std::string msg) {
return LDBCommandExecuteResult(EXEC_SUCCEED, msg);
}
static LDBCommandExecuteResult FAILED(std::string msg) {
return LDBCommandExecuteResult(EXEC_FAILED, msg);
}
private:
State state_;
std::string message_;
bool operator==(const LDBCommandExecuteResult&);
bool operator!=(const LDBCommandExecuteResult&);
};
class LDBCommand {
public:
/* Constructor */
LDBCommand(std::string& db_name, std::vector<std::string>& args) :
db_path_(db_name),
db_(NULL) {
parse_open_args(args);
}
LDBCommand(std::vector<std::string>& args) :
db_path_(""),
db_(NULL) {
parse_open_args(args);
}
virtual leveldb::Options PrepareOptionsForOpenDB();
virtual bool NoDBOpen() {
return false;
}
virtual ~LDBCommand() {
if (db_ != NULL) {
delete db_;
db_ = NULL;
}
}
/* Print the help message */
static void Help(std::string& ret) {
ret.append("--db=DB_PATH [");
ret.append(LDBCommand::BLOOM_ARG);
ret.append("<int,e.g.:14>] [");
ret.append(LDBCommand::COMPRESSION_TYPE_ARG);
ret.append("<no|snappy|zlib|bzip2> ");
ret.append(LDBCommand::BLOCK_SIZE);
ret.append("=<block_size_in_bytes> ");
ret.append(LDBCommand::AUTO_COMPACTION);
ret.append("=<true|false>]");
}
/* Run the command, and return the execute result. */
void Run() {
if (!exec_state_.IsNotStarted()) {
return;
}
if (db_ == NULL && !NoDBOpen()) {
OpenDB();
if (!exec_state_.IsNotStarted()) {
return;
}
}
DoCommand();
if (exec_state_.IsNotStarted()) {
exec_state_ = LDBCommandExecuteResult::SUCCEED("");
}
if (db_ != NULL) {
CloseDB ();
}
}
virtual void DoCommand() = 0;
LDBCommandExecuteResult GetExecuteState() {
return exec_state_;
}
void ClearPreviousRunState() {
exec_state_.Reset();
}
static std::string HexToString(const std::string& str) {
std::string parsed;
for (unsigned int i = 0; i < str.length();) {
int c;
sscanf(str.c_str() + i, "%2X", &c);
parsed.push_back(c);
i += 2;
}
return parsed;
}
protected:
void OpenDB() {
leveldb::Options opt = PrepareOptionsForOpenDB();
if (!exec_state_.IsNotStarted()) {
return;
}
// Open the DB.
leveldb::Status st = leveldb::DB::Open(opt, db_path_, &db_);
if (!st.ok()) {
std::string msg = st.ToString();
exec_state_ = LDBCommandExecuteResult::FAILED(msg);
}
}
void CloseDB () {
if (db_ != NULL) {
delete db_;
db_ = NULL;
}
}
static const char* FROM_ARG;
static const char* END_ARG;
static const char* HEX_ARG;
LDBCommandExecuteResult exec_state_;
std::string db_path_;
leveldb::DB* db_;
private:
static const char* BLOOM_ARG;
static const char* COMPRESSION_TYPE_ARG;
static const char* BLOCK_SIZE;
static const char* AUTO_COMPACTION;
std::vector<std::string> open_args_;
void parse_open_args(std::vector<std::string>& args);
};
class Compactor: public LDBCommand {
public:
Compactor(std::string& db_name, std::vector<std::string>& args);
virtual ~Compactor() {}
static void Help(std::string& ret);
virtual void DoCommand();
private:
bool null_from_;
std::string from_;
bool null_to_;
std::string to_;
bool hex_;
};
class DBDumper: public LDBCommand {
public:
DBDumper(std::string& db_name, std::vector<std::string>& args);
virtual ~DBDumper() {}
static void Help(std::string& ret);
virtual void DoCommand();
private:
bool null_from_;
std::string from_;
bool null_to_;
std::string to_;
int max_keys_;
bool count_only_;
bool print_stats_;
bool hex_;
bool hex_output_;
static const char* MAX_KEYS_ARG;
static const char* COUNT_ONLY_ARG;
static const char* STATS_ARG;
static const char* HEX_OUTPUT_ARG;
};
class DBLoader: public LDBCommand {
public:
DBLoader(std::string& db_name, std::vector<std::string>& args);
virtual ~DBLoader() {}
static void Help(std::string& ret);
virtual void DoCommand();
virtual leveldb::Options PrepareOptionsForOpenDB();
private:
bool hex_input_;
bool create_if_missing_;
bool disable_wal_;
static const char* HEX_INPUT_ARG;
static const char* CREATE_IF_MISSING_ARG;
static const char* DISABLE_WAL_ARG;
};
class ReduceDBLevels : public LDBCommand {
public:
ReduceDBLevels (std::string& db_name, std::vector<std::string>& args);
~ReduceDBLevels() {}
virtual leveldb::Options PrepareOptionsForOpenDB();
virtual void DoCommand();
virtual bool NoDBOpen() {
return true;
}
static void Help(std::string& msg);
static std::vector<std::string> PrepareArgs(int new_levels,
bool print_old_level = false);
private:
int old_levels_;
int new_levels_;
bool print_old_levels_;
int file_size_;
enum leveldb::CompressionType compression_;
static const char* NEW_LEVLES_ARG;
static const char* PRINT_OLD_LEVELS_ARG;
static const char* COMPRESSION_TYPE_ARG;
static const char* FILE_SIZE_ARG;
Status GetOldNumOfLevels(leveldb::Options& opt, int* levels);
};
class WALDumper : public LDBCommand {
public:
WALDumper (std::vector<std::string>& args);
~WALDumper() {}
virtual bool NoDBOpen() {
return true;
}
static void Help(std::string& ret);
virtual void DoCommand();
private:
bool print_header_;
std::string wal_file_;
static const char* WAL_FILE_ARG;
};
}
#endif