Merge pull request #456 from facebook/dbhandle-resiliancy

adding more resiliancy to DBHandle

include/osquery/database/db_handle.h

@@ -79,15 +79,6 @@ class DBHandle {
    */
   static std::shared_ptr<DBHandle> getInstance();
 
-  /**
-   * @brief Getter for the status of the operations required to open the
-   * database
-   *
-   * @return an instance of osquery::Status which indicates the success or
-   * failure of connecting to RocksDB
-   */
-  Status getStatus();
-
   /**
    * @brief Helper method which can be used to get a raw pointer to the
    * underlying RocksDB database handle
@@ -231,9 +222,6 @@ class DBHandle {
   /// The database handle
   rocksdb::DB* db_;
 
-  /// The status code that is generated while attempting to connect to RocksDB
-  rocksdb::Status status_;
-
   /// Column family descriptors which are used to connect to RocksDB
   std::vector<rocksdb::ColumnFamilyDescriptor> column_families_;
 

osquery/database/db_handle.cpp

@@ -28,7 +28,7 @@ const std::vector<std::string> kDomains = {kConfigurations, kQueries, kEvents};
 
 DEFINE_osquery_flag(string,
                     db_path,
-                    "/tmp/rocksdb-osquery",
+                    "/var/osquery/osquery.db",
                     "If using a disk-based backing store, specify a path.");
 
 DEFINE_osquery_flag(bool,
@@ -45,13 +45,13 @@ DBHandle::DBHandle(const std::string& path, bool in_memory) {
   options_.create_missing_column_families = true;
 
   if (in_memory) {
-    // Remove when upgrading to RocksDB 3.3
+    // Remove when MemEnv is included in librocksdb
     // options_.env = rocksdb::NewMemEnv(rocksdb::Env::Default());
-    throw std::domain_error("Required RocksDB 3.3 (and setMemEnv)");
+    throw std::runtime_error("Requires MemEnv");
   }
 
   if (pathExists(path).ok() && !isWritable(path).ok()) {
-    throw std::domain_error("Cannot write to RocksDB path: " + path);
+    throw std::runtime_error("Cannot write to RocksDB path: " + path);
   }
 
   column_families_.push_back(rocksdb::ColumnFamilyDescriptor(
@@ -62,8 +62,10 @@ DBHandle::DBHandle(const std::string& path, bool in_memory) {
         cf_name, rocksdb::ColumnFamilyOptions()));
   }
 
-  status_ =
-      rocksdb::DB::Open(options_, path, column_families_, &handles_, &db_);
+  auto s = rocksdb::DB::Open(options_, path, column_families_, &handles_, &db_);
+  if (!s.ok()) {
+    throw std::runtime_error(s.ToString());
+  }
 }
 
 DBHandle::~DBHandle() {
@@ -99,18 +101,18 @@ std::shared_ptr<DBHandle> DBHandle::getInstance(const std::string& path,
 // getters and setters
 /////////////////////////////////////////////////////////////////////////////
 
-osquery::Status DBHandle::getStatus() {
-  return Status(status_.code(), status_.ToString());
-}
-
 rocksdb::DB* DBHandle::getDB() { return db_; }
 
 rocksdb::ColumnFamilyHandle* DBHandle::getHandleForColumnFamily(
     const std::string& cf) {
-  for (int i = 0; i < kDomains.size(); i++) {
-    if (kDomains[i] == cf) {
-      return handles_[i];
+  try {
+    for (int i = 0; i < kDomains.size(); i++) {
+      if (kDomains[i] == cf) {
+        return handles_[i];
+      }
     }
+  } catch (const std::exception& e) {
+    // pass through and return nullptr
   }
   return nullptr;
 }
@@ -122,30 +124,45 @@ rocksdb::ColumnFamilyHandle* DBHandle::getHandleForColumnFamily(
 osquery::Status DBHandle::Get(const std::string& domain,
                               const std::string& key,
                               std::string& value) {
-  auto s = getDB()->Get(
-      rocksdb::ReadOptions(), getHandleForColumnFamily(domain), key, &value);
+  auto cfh = getHandleForColumnFamily(domain);
+  if (cfh == nullptr) {
+    return Status(1, "Could not get column family for " + domain);
+  }
+  auto s = getDB()->Get(rocksdb::ReadOptions(), cfh, key, &value);
   return Status(s.code(), s.ToString());
 }
 
 osquery::Status DBHandle::Put(const std::string& domain,
                               const std::string& key,
                               const std::string& value) {
-  auto s = getDB()->Put(
-      rocksdb::WriteOptions(), getHandleForColumnFamily(domain), key, value);
+  auto cfh = getHandleForColumnFamily(domain);
+  if (cfh == nullptr) {
+    return Status(1, "Could not get column family for " + domain);
+  }
+  auto s = getDB()->Put(rocksdb::WriteOptions(), cfh, key, value);
   return Status(s.code(), s.ToString());
 }
 
 osquery::Status DBHandle::Delete(const std::string& domain,
                                  const std::string& key) {
-  auto s = getDB()->Delete(
-      rocksdb::WriteOptions(), getHandleForColumnFamily(domain), key);
+  auto cfh = getHandleForColumnFamily(domain);
+  if (cfh == nullptr) {
+    return Status(1, "Could not get column family for " + domain);
+  }
+  auto s = getDB()->Delete(rocksdb::WriteOptions(), cfh, key);
   return Status(s.code(), s.ToString());
 }
 
 osquery::Status DBHandle::Scan(const std::string& domain,
                                std::vector<std::string>& results) {
-  auto it = getDB()->NewIterator(rocksdb::ReadOptions(),
-                                 getHandleForColumnFamily(domain));
+  auto cfh = getHandleForColumnFamily(domain);
+  if (cfh == nullptr) {
+    return Status(1, "Could not get column family for " + domain);
+  }
+  auto it = getDB()->NewIterator(rocksdb::ReadOptions(), cfh);
+  if (it == nullptr) {
+    return Status(1, "Could not get iterator for " + domain);
+  }
   for (it->SeekToFirst(); it->Valid(); it->Next()) {
     results.push_back(it->key().ToString());
   }

osquery/database/db_handle_tests.cpp

@@ -31,18 +31,9 @@ class DBHandleTests : public testing::Test {
   std::shared_ptr<DBHandle> db;
 };
 
-TEST_F(DBHandleTests, test_create_new_database_on_disk) {
-  EXPECT_TRUE(db->getStatus().ok());
-  EXPECT_EQ(db->getStatus().toString(), "OK");
-}
-
 TEST_F(DBHandleTests, test_singleton_on_disk) {
   auto db1 = DBHandle::getInstance();
-  EXPECT_TRUE(db1->getStatus().ok());
-  EXPECT_EQ(db1->getStatus().toString(), "OK");
   auto db2 = DBHandle::getInstance();
-  EXPECT_TRUE(db2->getStatus().ok());
-  EXPECT_EQ(db2->getStatus().toString(), "OK");
   EXPECT_EQ(db1, db2);
 }