From be9b862d473de3868bdd2efaff889928bb47c6b8 Mon Sep 17 00:00:00 2001
From: Zheng Shao <zshao@fb.com>
Date: Sun, 16 Dec 2012 17:06:51 -0800
Subject: [PATCH] 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
---
 tools/ldb.cc    |  8 +++++
 util/ldb_cmd.cc | 84 +++++++++++++++++++++++++++++++++++++++++++++++--
 util/ldb_cmd.h  | 19 +++++++++++
 3 files changed, 109 insertions(+), 2 deletions(-)

diff --git a/tools/ldb.cc b/tools/ldb.cc
index c4d7c714b..86abd468c 100644
--- a/tools/ldb.cc
+++ b/tools/ldb.cc
@@ -21,6 +21,11 @@ public:
     ret.append(" dump ");
     DBDumper::Help(ret);
 
+    ret.append("\n--- load ----:\n");
+    ret.append(exec_name);
+    ret.append(" load ");
+    DBLoader::Help(ret);
+
     ret.append("\n---reduce_levels ----:\n");
     ret.append(exec_name);
     ret.append(" reduce_levels ");
@@ -56,6 +61,9 @@ public:
     } else if (strcmp(cmd, "dump") == 0) {
       // run dump
       cmdObj = new DBDumper(db_name, args);
+    } else if (strcmp(cmd, "load") == 0) {
+      // run loader
+      cmdObj = new DBLoader(db_name, args);
     } else if (strcmp(cmd, "reduce_levels") == 0) {
       // reduce db levels
       cmdObj = new ReduceDBLevels(db_name, args);
diff --git a/util/ldb_cmd.cc b/util/ldb_cmd.cc
index 23c259e07..9d1f703d5 100644
--- a/util/ldb_cmd.cc
+++ b/util/ldb_cmd.cc
@@ -144,6 +144,85 @@ void Compactor::DoCommand() {
   delete end;
 }
 
+const char* DBLoader::HEX_INPUT_ARG = "--input_hex";
+const char* DBLoader::CREATE_IF_MISSING_ARG = "--create_if_missing";
+const char* DBLoader::DISABLE_WAL_ARG = "--disable_wal";
+static const char* delim = " ==> ";
+
+DBLoader::DBLoader(std::string& db_name, std::vector<std::string>& args) :
+    LDBCommand(db_name, args),
+    hex_input_(false),
+    create_if_missing_(false) {
+  for (unsigned int i = 0; i < args.size(); i++) {
+    std::string& arg = args.at(i);
+    if (arg.find(HEX_INPUT_ARG) == 0) {
+      hex_input_ = true;
+    } else if (arg.find(CREATE_IF_MISSING_ARG) == 0) {
+      create_if_missing_ = true;
+    } else if (arg.find(DISABLE_WAL_ARG) == 0) {
+      disable_wal_ = true;
+    } else {
+      exec_state_ = LDBCommandExecuteResult::FAILED("Unknown argument:" + arg);
+    }
+  }
+}
+
+void DBLoader::Help(std::string& ret) {
+  LDBCommand::Help(ret);
+  ret.append("[");
+  ret.append(HEX_INPUT_ARG);
+  ret.append("] [");
+  ret.append(CREATE_IF_MISSING_ARG);
+  ret.append("] [");
+  ret.append(DISABLE_WAL_ARG);
+  ret.append("]");
+}
+
+leveldb::Options DBLoader::PrepareOptionsForOpenDB() {
+  leveldb::Options opt = LDBCommand::PrepareOptionsForOpenDB();
+  opt.create_if_missing = create_if_missing_;
+  return opt;
+}
+
+void DBLoader::DoCommand() {
+  if (!db_) {
+    return;
+  }
+
+  WriteOptions write_options;
+  if (disable_wal_) {
+    write_options.disableWAL = true;
+  }
+
+  int bad_lines = 0;
+  std::string line;
+  while (std::getline(std::cin, line, '\n')) {
+    size_t pos = line.find(delim);
+    if (pos != std::string::npos) {
+      std::string key = line.substr(0, pos);
+      std::string value = line.substr(pos + strlen(delim));
+
+      if (hex_input_) {
+        key = HexToString(key);
+        value = HexToString(value);
+      }
+
+      db_->Put(write_options, Slice(key), Slice(value));
+
+    } else if (0 == line.find("Keys in range:")) {
+      // ignore this line
+    } else if (0 == line.find("Created bg thread 0x")) {
+      // ignore this line
+    } else {
+      bad_lines ++;
+    }
+  }
+  
+  if (bad_lines > 0) {
+    std::cout << "Warning: " << bad_lines << " bad lines ignored." << std::endl;
+  }
+}
+
 const char* DBDumper::MAX_KEYS_ARG = "--max_keys=";
 const char* DBDumper::COUNT_ONLY_ARG = "--count_only";
 const char* DBDumper::STATS_ARG = "--stats";
@@ -247,14 +326,15 @@ void DBDumper::DoCommand() {
         for (unsigned int i = 0; i < str.length(); ++i) {
           fprintf(stdout, "%02X", (unsigned char)str[i]);
         }
-        fprintf(stdout, " ==> ");
+        fprintf(stdout, delim);
         str = iter->value().ToString();
         for (unsigned int i = 0; i < str.length(); ++i) {
           fprintf(stdout, "%02X", (unsigned char)str[i]);
         }
         fprintf(stdout, "\n");
       } else {
-        fprintf(stdout, "%s ==> %s\n", iter->key().ToString().c_str(),
+        fprintf(stdout, "%s%s%s\n", iter->key().ToString().c_str(),
+            delim,
             iter->value().ToString().c_str());
       }
     }
diff --git a/util/ldb_cmd.h b/util/ldb_cmd.h
index ae4351046..d061535c8 100644
--- a/util/ldb_cmd.h
+++ b/util/ldb_cmd.h
@@ -254,6 +254,25 @@ private:
   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: