Implement DBCloud::ListColumnFamilies()

Summary: We need to load cloud manifest before calling into DB::ListColumnFamilies().

Test Plan: Added a unit test

Reviewers: #sys-eng_team, dhruba

Differential Revision: https://rockset.phacility.com/D11081

Author: igorcanadi

cloud/db_cloud_impl.cc

@@ -120,13 +120,17 @@ Status DBCloud::Open(const Options& opt, const std::string& local_dbname,
         local_dbname);  // MJR: TODO: Move into sanitize
   }
 
-  st = cenv->SanitizeDirectory(options, local_dbname, read_only);
+  // If cloud manifest is already loaded, this means the directory has been
+  // sanitized (possibly by the call to ListColumnFamilies())
+  if (cenv->GetCloudManifest() == nullptr) {
+    st = cenv->SanitizeDirectory(options, local_dbname, read_only);
 
-  if (st.ok()) {
-    st = cenv->LoadCloudManifest(local_dbname, read_only);
-  }
-  if (!st.ok()) {
-    return st;
+    if (st.ok()) {
+      st = cenv->LoadCloudManifest(local_dbname, read_only);
+    }
+    if (!st.ok()) {
+      return st;
+    }
   }
   // If a persistent cache path is specified, then we set it in the options.
   if (!persistent_cache_path.empty() && persistent_cache_size_gb) {
@@ -403,5 +407,25 @@ Status DBCloudImpl::ExecuteRemoteCompactionRequest(
   return Status::OK();
 }
 
+Status DBCloud::ListColumnFamilies(const DBOptions& db_options,
+                                   const std::string& name,
+                                   std::vector<std::string>* column_families) {
+  CloudEnvImpl* cenv = static_cast<CloudEnvImpl*>(db_options.env);
+
+  Env* local_env = cenv->GetBaseEnv();
+  local_env->CreateDirIfMissing(name);
+
+  Status st;
+  st = cenv->SanitizeDirectory(db_options, name, false);
+  if (st.ok()) {
+    st = cenv->LoadCloudManifest(name, false);
+  }
+  if (st.ok()) {
+    st = DB::ListColumnFamilies(db_options, name, column_families);
+  }
+
+  return st;
+}
+
 }  // namespace ROCKSDB_NAMESPACE
 #endif  // ROCKSDB_LITE

cloud/db_cloud_test.cc

@@ -130,6 +130,24 @@ class CloudTest : public testing::Test {
 
   // Open database via the cloud interface
   void OpenDB() {
+    // default column family
+    ColumnFamilyOptions cfopt = options_;
+    std::vector<ColumnFamilyDescriptor> column_families;
+    column_families.emplace_back(
+        ColumnFamilyDescriptor(kDefaultColumnFamilyName, cfopt));
+
+    OpenWithColumnFamilies({kDefaultColumnFamilyName});
+  }
+
+  void CreateColumnFamilies(const std::vector<std::string>& cfs) {
+    size_t cfi = handles_.size();
+    handles_.resize(cfi + cfs.size());
+    for (auto cf : cfs) {
+      ASSERT_OK(db_->CreateColumnFamily(options_, cf, &handles_[cfi++]));
+    }
+  }
+
+  void OpenWithColumnFamilies(const std::vector<std::string>& cfs) {
     ASSERT_TRUE(cloud_env_options_.credentials.HasValid().ok());
 
     // Create new AWS env
@@ -139,23 +157,15 @@ class CloudTest : public testing::Test {
     // Sleep for a second because S3 is eventual consistency.
     std::this_thread::sleep_for(std::chrono::seconds(1));
 
-    // default column family
-    ColumnFamilyOptions cfopt = options_;
-    std::vector<ColumnFamilyDescriptor> column_families;
-    column_families.emplace_back(
-        ColumnFamilyDescriptor(kDefaultColumnFamilyName, cfopt));
-    std::vector<ColumnFamilyHandle*> handles;
-
     ASSERT_TRUE(db_ == nullptr);
+    std::vector<ColumnFamilyDescriptor> column_families;
+    for (size_t i = 0; i < cfs.size(); ++i) {
+      column_families.emplace_back(cfs[i], options_);
+    }
     ASSERT_OK(DBCloud::Open(options_, dbname_, column_families,
                             persistent_cache_path_, persistent_cache_size_gb_,
-                            &handles, &db_));
+                            &handles_, &db_));
     ASSERT_OK(db_->GetDbIdentity(dbid_));
-
-    // Delete the handle for the default column family because the DBImpl
-    // always holds a reference to it.
-    ASSERT_TRUE(handles.size() > 0);
-    delete handles[0];
   }
 
   // Creates and Opens a clone
@@ -224,6 +234,10 @@ class CloudTest : public testing::Test {
   }
 
   void CloseDB() {
+    for (auto h : handles_) {
+      delete h;
+    }
+    handles_.clear();
     if (db_) {
       db_->Flush(FlushOptions());  // convert pending writes to sst files
       delete db_;
@@ -288,6 +302,8 @@ class CloudTest : public testing::Test {
   uint64_t persistent_cache_size_gb_;
   DBCloud* db_;
   std::unique_ptr<CloudEnv> aenv_;
+
+  std::vector<ColumnFamilyHandle*> handles_;
 };
 
 //
@@ -422,6 +438,53 @@ TEST_F(CloudTest, Newdb) {
   CloseDB();
 }
 
+TEST_F(CloudTest, ColumnFamilies) {
+  // Put one key-value
+  OpenDB();
+
+  CreateColumnFamilies({"cf1", "cf2"});
+
+  ASSERT_OK(db_->Put(WriteOptions(), handles_[0], "hello", "a"));
+  ASSERT_OK(db_->Put(WriteOptions(), handles_[1], "hello", "b"));
+  ASSERT_OK(db_->Put(WriteOptions(), handles_[2], "hello", "c"));
+
+  auto validate = [&]() {
+    std::string value;
+    ASSERT_OK(db_->Get(ReadOptions(), handles_[0], "hello", &value));
+    ASSERT_EQ(value, "a");
+    ASSERT_OK(db_->Get(ReadOptions(), handles_[1], "hello", &value));
+    ASSERT_EQ(value, "b");
+    ASSERT_OK(db_->Get(ReadOptions(), handles_[2], "hello", &value));
+    ASSERT_EQ(value, "c");
+  };
+
+  validate();
+
+  CloseDB();
+  OpenWithColumnFamilies({kDefaultColumnFamilyName, "cf1", "cf2"});
+
+  validate();
+  CloseDB();
+
+  // destory local state
+  DestroyDir(dbname_);
+
+  // new aws env
+  CreateAwsEnv();
+  options_.env = aenv_.get();
+
+  std::vector<std::string> families;
+  ASSERT_OK(DBCloud::ListColumnFamilies(options_, dbname_, &families));
+  std::sort(families.begin(), families.end());
+  ASSERT_TRUE(families == std::vector<std::string>(
+                              {"cf1", "cf2", kDefaultColumnFamilyName}));
+
+
+  OpenWithColumnFamilies({kDefaultColumnFamilyName, "cf1", "cf2"});
+  validate();
+  CloseDB();
+}
+
 //
 // Create and read from a clone.
 //

include/rocksdb/cloud/db_cloud.h

@@ -62,6 +62,14 @@ class DBCloud : public StackableDB {
       const PluggableCompactionParam& inputParams,
       PluggableCompactionResult* result, bool sanitize) = 0;
 
+  // ListColumnFamilies will open the DB specified by argument name
+  // and return the list of all column families in that DB
+  // through column_families argument. The ordering of
+  // column families in column_families is unspecified.
+  static Status ListColumnFamilies(const DBOptions& db_options,
+                                   const std::string& name,
+                                   std::vector<std::string>* column_families);
+
   virtual ~DBCloud() {}
 
  protected: