Last active
November 14, 2022 11:37
-
-
Save liukai/9379499 to your computer and use it in GitHub Desktop.
Migrate RocksDB database with different options
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| rocksdb::Status CopyDB(const std::string& source_db_name, | |
| const Options& source_db_options, | |
| const std::string& destination_db_name, | |
| const Options& destination_db_options) { | |
| rocksdb::Status s; | |
| // Create the source/destination DBs | |
| rocksdb::DB* source_db = nullptr; | |
| auto sanitized_source_db_options = source_db_options; | |
| sanitized_source_db_options.create_if_missing = false; | |
| s = rocksdb::DB::Open(sanitized_source_db_options, source_db_name, &source_db); | |
| std::unique_ptr<rocksdb::DB> source_db_guard(source_db); | |
| if (!s.ok()) { | |
| return s; | |
| } | |
| rocksdb::DB* destination_db = nullptr; | |
| auto sanitized_destination_db_options = destination_db_options; | |
| sanitized_destination_db_options.error_if_exists = true; | |
| sanitized_destination_db_options.create_if_missing = true; | |
| s = rocksdb::DB::Open(sanitized_destination_db_options, | |
| destination_db_name, | |
| &destination_db); | |
| std::unique_ptr<rocksdb::DB> destination_db_guard(destination_db); | |
| if (!s.ok()) { | |
| return s; | |
| } | |
| std::unique_ptr<rocksdb::Iterator> source_iterator( | |
| source_db->NewIterator(rocksdb::ReadOptions())); | |
| WriteOptions write_options; | |
| for (source_iterator->SeekToFirst(); | |
| source_iterator->Valid() && source_iterator->status().ok(); | |
| source_iterator->Next()) { | |
| destination_db->Put(write_options, source_iterator->key(), | |
| source_iterator->value()); | |
| } | |
| return s; | |
| } | |
| // TEST | |
| std::vector<Options> OptionsList() { | |
| std::vector<Options> options_list; | |
| { | |
| Options options; | |
| options_list.push_back(options); | |
| } | |
| { | |
| Options options; | |
| options.memtable_factory.reset( | |
| NewHashSkipListRepFactory(NewFixedPrefixTransform(1))); | |
| options_list.push_back(options); | |
| } | |
| { | |
| Options options; | |
| options.memtable_factory.reset( | |
| NewHashLinkListRepFactory(NewFixedPrefixTransform(1), 4)); | |
| options_list.push_back(options); | |
| } | |
| { | |
| Options options; | |
| options.filter_policy = NewBloomFilterPolicy(10); | |
| options_list.push_back(options); | |
| } | |
| return options_list; | |
| } | |
| TEST(DBTest, CopyDBTest) { | |
| Random rnd(1028); | |
| auto source = test::TmpDir() + "/" + "db_source"; | |
| auto destination = test::TmpDir() + "/" + "db_destination"; | |
| DestroyDB(source, Options()); | |
| // Create a db with default options | |
| Options src_opt; | |
| src_opt.create_if_missing = true; | |
| DB* db_from = nullptr; | |
| auto s = DB::Open(src_opt, source, &db_from); | |
| if (!s.ok()) { | |
| delete db_from; | |
| } | |
| ASSERT_OK(s); | |
| const int kEntrySize = 10000; | |
| for (int i = 0; i < kEntrySize; ++i) { | |
| db_from->Put(WriteOptions(), Key(i), RandomString(&rnd, 1500)); | |
| } | |
| delete db_from; // release db and flush it. | |
| // copy data | |
| for (const auto& dest_opt : OptionsList()) { | |
| DestroyDB(destination, dest_opt); | |
| Options dest_options; | |
| s = CopyDB(source, src_opt, destination, dest_opt); | |
| ASSERT_OK(s); | |
| // compare data | |
| s = DB::Open(src_opt, source, &db_from); | |
| std::unique_ptr<DB> src_db_guard(db_from); | |
| ASSERT_OK(s); | |
| DB* db_to; | |
| s = DB::Open(dest_opt, destination, &db_to); | |
| std::unique_ptr<DB> dest_db_guard(db_to); | |
| ASSERT_OK(s); | |
| std::unique_ptr<Iterator> source_iterator( | |
| db_from->NewIterator(ReadOptions())); | |
| std::unique_ptr<Iterator> dest_iterator(db_to->NewIterator(ReadOptions())); | |
| for (source_iterator->SeekToFirst(), dest_iterator->SeekToFirst(); | |
| source_iterator->Valid() && source_iterator->status().ok() && | |
| dest_iterator->Valid() && dest_iterator->status().ok(); | |
| source_iterator->Next(), dest_iterator->Next()) { | |
| ASSERT_EQ(source_iterator->key().compare(dest_iterator->key()), 0); | |
| ASSERT_EQ(source_iterator->value().compare(dest_iterator->value()), 0); | |
| } | |
| ASSERT_OK(source_iterator->status()); | |
| ASSERT_OK(dest_iterator->status()); | |
| } | |
| } | |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment