Enhance ldb_cmd_tool to enable user pass in customized cfds (#12261)

Summary:
The current implementation of the ldb_cmd tool involves commenting out the user-passed column_family_descriptors, resulting in the tool consistently constructing its column_family_descriptors from the pre-existing OPTIONS file.

The proposed fix prioritizes user-passed column family descriptors, ensuring they take precedence over those specified in the OPTIONS file. This modification enhances the tool's adaptability and responsiveness to user configurations.

Pull Request resolved: #12261

Reviewed By: cbi42

Differential Revision: D52965877

Pulled By: ajkr

fbshipit-source-id: 334a83a8e1004c271b19e7ca09381a0e7cf87b03

tools/ldb_cmd.cc

@@ -150,7 +150,7 @@ LDBCommand* LDBCommand::InitFromCmdLineArgs(
 LDBCommand* LDBCommand::InitFromCmdLineArgs(
     const std::vector<std::string>& args, const Options& options,
     const LDBOptions& ldb_options,
-    const std::vector<ColumnFamilyDescriptor>* /*column_families*/,
+    const std::vector<ColumnFamilyDescriptor>* column_families,
     const std::function<LDBCommand*(const ParsedParams&)>& selector) {
   // --x=y command line arguments are added as x->y map entries in
   // parsed_params.option_map.
@@ -200,6 +200,7 @@ LDBCommand* LDBCommand::InitFromCmdLineArgs(
   if (command) {
     command->SetDBOptions(options);
     command->SetLDBOptions(ldb_options);
+    command->SetColumnFamilies(column_families);
   }
   return command;
 }
@@ -933,10 +934,12 @@ void LDBCommand::OverrideBaseCFOptions(ColumnFamilyOptions* cf_opts) {
 // Second, overrides the options according to the CLI arguments and the
 // specific subcommand being run.
 void LDBCommand::PrepareOptions() {
+  std::vector<ColumnFamilyDescriptor> column_families_from_options;
+
   if (!create_if_missing_ && try_load_options_) {
     config_options_.env = options_.env;
     Status s = LoadLatestOptions(config_options_, db_path_, &options_,
-                                 &column_families_);
+                                 &column_families_from_options);
     if (!s.ok() && !s.IsNotFound()) {
       // Option file exists but load option file error.
       std::string current_version = std::to_string(ROCKSDB_MAJOR) + "." +
@@ -962,7 +965,7 @@ void LDBCommand::PrepareOptions() {
     }
 
     // If merge operator is not set, set a string append operator.
-    for (auto& cf_entry : column_families_) {
+    for (auto& cf_entry : column_families_from_options) {
       if (!cf_entry.options.merge_operator) {
         cf_entry.options.merge_operator =
             MergeOperators::CreateStringAppendOperator(':');
@@ -980,22 +983,29 @@ void LDBCommand::PrepareOptions() {
   }
 
   if (column_families_.empty()) {
-    // Reads the MANIFEST to figure out what column families exist. In this
-    // case, the option overrides from the CLI argument/specific subcommand
-    // apply to all column families.
-    std::vector<std::string> cf_list;
-    Status st = DB::ListColumnFamilies(options_, db_path_, &cf_list);
-    // It is possible the DB doesn't exist yet, for "create if not
-    // existing" case. The failure is ignored here. We rely on DB::Open()
-    // to give us the correct error message for problem with opening
-    // existing DB.
-    if (st.ok() && cf_list.size() > 1) {
-      // Ignore single column family DB.
-      for (const auto& cf_name : cf_list) {
-        column_families_.emplace_back(cf_name, options_);
+    // column_families not set. Either set it from MANIFEST or OPTIONS file.
+    if (column_families_from_options.empty()) {
+      // Reads the MANIFEST to figure out what column families exist. In this
+      // case, the option overrides from the CLI argument/specific subcommand
+      // apply to all column families.
+      std::vector<std::string> cf_list;
+      Status st = DB::ListColumnFamilies(options_, db_path_, &cf_list);
+      // It is possible the DB doesn't exist yet, for "create if not
+      // existing" case. The failure is ignored here. We rely on DB::Open()
+      // to give us the correct error message for problem with opening
+      // existing DB.
+      if (st.ok() && cf_list.size() > 1) {
+        // Ignore single column family DB.
+        for (const auto& cf_name : cf_list) {
+          column_families_.emplace_back(cf_name, options_);
+        }
       }
+    } else {
+      SetColumnFamilies(&column_families_from_options);
     }
-  } else {
+  }
+
+  if (!column_families_from_options.empty()) {
     // We got column families from the OPTIONS file. In this case, the option
     // overrides from the CLI argument/specific subcommand only apply to the
     // column family specified by `--column_family_name`.

tools/ldb_cmd_test.cc

@@ -1232,7 +1232,8 @@ TEST_F(LdbCmdTest, CustomComparator) {
   std::string dbname = test::PerThreadDBPath(env, "ldb_cmd_test");
   DB* db = nullptr;
 
-  std::vector<ColumnFamilyDescriptor> cfds = {{kDefaultColumnFamilyName, opts}};
+  std::vector<ColumnFamilyDescriptor> cfds = {
+      {kDefaultColumnFamilyName, opts}, {"cf1", opts}, {"cf2", opts}};
   std::vector<ColumnFamilyHandle*> handles;
   ASSERT_OK(DestroyDB(dbname, opts));
   ASSERT_OK(DB::Open(opts, dbname, cfds, &handles, &db));
@@ -1250,7 +1251,7 @@ TEST_F(LdbCmdTest, CustomComparator) {
   char* argv[] = {arg1, const_cast<char*>(arg2.c_str()), arg3, arg4};
 
   ASSERT_EQ(0,
-            LDBCommandRunner::RunCommand(4, argv, opts, LDBOptions(), nullptr));
+            LDBCommandRunner::RunCommand(4, argv, opts, LDBOptions(), &cfds));
 }
 
 }  // namespace ROCKSDB_NAMESPACE