/Users/deen/code/yugabyte-db/src/yb/rocksdb/tools/dump/db_dump_tool.cc

Source (jump to first uncovered line)
//  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
//  This source code is licensed under the BSD-style license found in the
//  LICENSE file in the root directory of this source tree. An additional grant
//  of patent rights can be found in the PATENTS file in the same directory.
//
// The following only applies to changes made to this file as part of YugaByte development.
//
// Portions Copyright (c) YugaByte, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
// in compliance with the License.  You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software distributed under the License
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
// or implied.  See the License for the specific language governing permissions and limitations
// under the License.
//

#ifndef ROCKSDB_LITE

#ifndef __STDC_FORMAT_MACROS
#define __STDC_FORMAT_MACROS
#endif

#include <inttypes.h>
#include <iostream>

#include "yb/rocksdb/db.h"
#include "yb/rocksdb/db_dump_tool.h"
#include "yb/rocksdb/env.h"
#include "yb/rocksdb/util/coding.h"

namespace rocksdb {

bool DbDumpTool::Run(const DumpOptions& dump_options,
                     rocksdb::Options options) {
  rocksdb::DB* dbptr;
  rocksdb::Status status;
  std::unique_ptr<rocksdb::WritableFile> dumpfile;
  char hostname[1024];
  int64_t timesec;
  std::string abspath;
  char json[4096];

  static const char* magicstr = "ROCKDUMP";
  static const char versionstr[8] = {0, 0, 0, 0, 0, 0, 0, 1};

  rocksdb::Env* env = rocksdb::Env::Default();

  // Open the database
  options.create_if_missing = false;
  status = rocksdb::DB::OpenForReadOnly(options, dump_options.db_path, &dbptr);
  if (!status.ok()) {
    std::cerr << "Unable to open database '" << dump_options.db_path
              << "' for reading: " << status.ToString() << std::endl;
    return false;
  }

  const std::unique_ptr<rocksdb::DB> db(dbptr);

  status = env->NewWritableFile(dump_options.dump_location, &dumpfile,
                                rocksdb::EnvOptions());
  if (!status.ok()) {
    std::cerr << "Unable to open dump file '" << dump_options.dump_location
              << "' for writing: " << status.ToString() << std::endl;
    return false;
  }

  rocksdb::Slice magicslice(magicstr, 8);
  status = dumpfile->Append(magicslice);
  if (!status.ok()) {
    std::cerr << "Append failed: " << status.ToString() << std::endl;
    return false;
  }

  rocksdb::Slice versionslice(versionstr, 8);
  status = dumpfile->Append(versionslice);
  if (!status.ok()) {
    std::cerr << "Append failed: " << status.ToString() << std::endl;
    return false;
  }

  if (dump_options.anonymous) {
    snprintf(json, sizeof(json), "{}");
  } else {
    status = env->GetHostName(hostname, sizeof(hostname));
    status = env->GetCurrentTime(&timesec);
    status = env->GetAbsolutePath(dump_options.db_path, &abspath);
    snprintf(json, sizeof(json),
             "{ \"database-path\": \"%s\", \"hostname\": \"%s\", "
             "\"creation-time\": %" PRIi64 " }",
             abspath.c_str(), hostname, timesec);
  }

  rocksdb::Slice infoslice(json, strlen(json));
  char infosize[4];
  rocksdb::EncodeFixed32(infosize, (uint32_t)infoslice.size());
  rocksdb::Slice infosizeslice(infosize, 4);
  status = dumpfile->Append(infosizeslice);
  if (!status.ok()) {
    std::cerr << "Append failed: " << status.ToString() << std::endl;
    return false;
  }
  status = dumpfile->Append(infoslice);
  if (!status.ok()) {
    std::cerr << "Append failed: " << status.ToString() << std::endl;
    return false;
  }

  const std::unique_ptr<rocksdb::Iterator> it(
      db->NewIterator(rocksdb::ReadOptions()));
  for (it->SeekToFirst(); it->Valid(); it->Next()) {
    char keysize[4];
    rocksdb::EncodeFixed32(keysize, (uint32_t)it->key().size());
    rocksdb::Slice keysizeslice(keysize, 4);
    status = dumpfile->Append(keysizeslice);
    if (!status.ok()) {
      std::cerr << "Append failed: " << status.ToString() << std::endl;
      return false;
    }
    status = dumpfile->Append(it->key());
    if (!status.ok()) {
      std::cerr << "Append failed: " << status.ToString() << std::endl;
      return false;
    }

    char valsize[4];
    rocksdb::EncodeFixed32(valsize, (uint32_t)it->value().size());
    rocksdb::Slice valsizeslice(valsize, 4);
    status = dumpfile->Append(valsizeslice);
    if (!status.ok()) {
      std::cerr << "Append failed: " << status.ToString() << std::endl;
      return false;
    }
    status = dumpfile->Append(it->value());
    if (!status.ok()) {
      std::cerr << "Append failed: " << status.ToString() << std::endl;
      return false;
    }
  }
  if (!it->status().ok()) {
    std::cerr << "Database iteration failed: " << status.ToString()
              << std::endl;
    return false;
  }
  return true;
}

bool DbUndumpTool::Run(const UndumpOptions& undump_options,
                       rocksdb::Options options) {
  rocksdb::DB* dbptr;
  rocksdb::Status status;
  rocksdb::Env* env;
  std::unique_ptr<rocksdb::SequentialFile> dumpfile;
  rocksdb::Slice slice;
  uint8_t scratch8[8];

  static const char* magicstr = "ROCKDUMP";
  static const char versionstr[8] = {0, 0, 0, 0, 0, 0, 0, 1};

  env = rocksdb::Env::Default();

  status = env->NewSequentialFile(undump_options.dump_location, &dumpfile,
                                  rocksdb::EnvOptions());
  if (!status.ok()) {
    std::cerr << "Unable to open dump file '" << undump_options.dump_location
              << "' for reading: " << status.ToString() << std::endl;
    return false;
  }

  status = dumpfile->Read(8, &slice, scratch8);
  if (!status.ok() || slice.size() != 8 ||
      memcmp(slice.data(), magicstr, 8) != 0) {
    std::cerr << "File '" << undump_options.dump_location
              << "' is not a recognizable dump file." << std::endl;
    return false;
  }

  status = dumpfile->Read(8, &slice, scratch8);
  if (!status.ok() || slice.size() != 8 ||
      memcmp(slice.data(), versionstr, 8) != 0) {
    std::cerr << "File '" << undump_options.dump_location
              << "' version not recognized." << std::endl;
    return false;
  }

  status = dumpfile->Read(4, &slice, scratch8);
  if (!status.ok() || slice.size() != 4) {
    std::cerr << "Unable to read info blob size." << std::endl;
    return false;
  }
  uint32_t infosize = rocksdb::DecodeFixed32(slice.data());
  status = dumpfile->Skip(infosize);
  if (!status.ok()) {
    std::cerr << "Unable to skip info blob: " << status.ToString() << std::endl;
    return false;
  }

  options.create_if_missing = true;
  status = rocksdb::DB::Open(options, undump_options.db_path, &dbptr);
  if (!status.ok()) {
    std::cerr << "Unable to open database '" << undump_options.db_path
              << "' for writing: " << status.ToString() << std::endl;
    return false;
  }

  const std::unique_ptr<rocksdb::DB> db(dbptr);

  uint32_t last_keysize = 64;
  size_t last_valsize = 1 << 20;
  std::unique_ptr<uint8_t[]> keyscratch(new uint8_t[last_keysize]);
  std::unique_ptr<uint8_t[]> valscratch(new uint8_t[last_valsize]);

  while (1) {
    uint32_t keysize, valsize;
    rocksdb::Slice keyslice;
    rocksdb::Slice valslice;

    status = dumpfile->Read(4, &slice, scratch8);
    if (!status.ok() || slice.size() != 4) break;
    keysize = rocksdb::DecodeFixed32(slice.data());
    if (keysize > last_keysize) {
      while (keysize > last_keysize) last_keysize *= 2;
      keyscratch = std::unique_ptr<uint8_t[]>(new uint8_t[last_keysize]);
    }

    status = dumpfile->Read(keysize, &keyslice, keyscratch.get());
    if (!status.ok() || keyslice.size() != keysize) {
      std::cerr << "Key read failure: "
                << (status.ok() ? "insufficient data" : status.ToString())
                << std::endl;
      return false;
    }

    status = dumpfile->Read(4, &slice, scratch8);
    if (!status.ok() || slice.size() != 4) {
      std::cerr << "Unable to read value size: "
                << (status.ok() ? "insufficient data" : status.ToString())
                << std::endl;
      return false;
    }
    valsize = rocksdb::DecodeFixed32(slice.data());
    if (valsize > last_valsize) {
      while (valsize > last_valsize) last_valsize *= 2;
      valscratch = std::unique_ptr<uint8_t[]>(new uint8_t[last_valsize]);
    }

    status = dumpfile->Read(valsize, &valslice, valscratch.get());
    if (!status.ok() || valslice.size() != valsize) {
      std::cerr << "Unable to read value: "
                << (status.ok() ? "insufficient data" : status.ToString())
                << std::endl;
      return false;
    }

    status = db->Put(rocksdb::WriteOptions(), keyslice, valslice);
    if (!status.ok()) {
      fprintf(stderr, "Unable to write database entry\n");
      return false;
    }
  }

  if (undump_options.compact_db) {
    status = db->CompactRange(rocksdb::CompactRangeOptions(), nullptr, nullptr);
    if (!status.ok()) {
      fprintf(stderr,
              "Unable to compact the database after loading the dumped file\n");
      return false;
    }
  }
  return true;
}
}  // namespace rocksdb
#endif  // ROCKSDB_LITE

YugabyteDB (2.13.1.0-b60, 21121d69985fbf76aa6958d8f04a9bfa936293b5)

Coverage Report

Created: 2022-03-22 16:43

Line	Count	Source (jump to first uncovered line)
1		// Copyright (c) 2011-present, Facebook, Inc. All rights reserved.
2		// This source code is licensed under the BSD-style license found in the
3		// LICENSE file in the root directory of this source tree. An additional grant
4		// of patent rights can be found in the PATENTS file in the same directory.
5		//
6		// The following only applies to changes made to this file as part of YugaByte development.
7		//
8		// Portions Copyright (c) YugaByte, Inc.
9		//
10		// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
11		// in compliance with the License. You may obtain a copy of the License at
12		//
13		// http://www.apache.org/licenses/LICENSE-2.0
14		//
15		// Unless required by applicable law or agreed to in writing, software distributed under the License
16		// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
17		// or implied. See the License for the specific language governing permissions and limitations
18		// under the License.
19		//
20
21		#ifndef ROCKSDB_LITE
22
23		#ifndef __STDC_FORMAT_MACROS
24		#define __STDC_FORMAT_MACROS
25		#endif
26
27		#include <inttypes.h>
28		#include <iostream>
29
30		#include "yb/rocksdb/db.h"
31		#include "yb/rocksdb/db_dump_tool.h"
32		#include "yb/rocksdb/env.h"
33		#include "yb/rocksdb/util/coding.h"
34
35		namespace rocksdb {
36
37		bool DbDumpTool::Run(const DumpOptions& dump_options,
38	0	rocksdb::Options options) {
39	0	rocksdb::DB* dbptr;
40	0	rocksdb::Status status;
41	0	std::unique_ptr<rocksdb::WritableFile> dumpfile;
42	0	char hostname[1024];
43	0	int64_t timesec;
44	0	std::string abspath;
45	0	char json[4096];
46
47	0	static const char* magicstr = "ROCKDUMP";
48	0	static const char versionstr[8] = {0, 0, 0, 0, 0, 0, 0, 1};
49
50	0	rocksdb::Env* env = rocksdb::Env::Default();
51
52		// Open the database
53	0	options.create_if_missing = false;
54	0	status = rocksdb::DB::OpenForReadOnly(options, dump_options.db_path, &dbptr);
55	0	if (!status.ok()) {
56	0	std::cerr << "Unable to open database '" << dump_options.db_path
57	0	<< "' for reading: " << status.ToString() << std::endl;
58	0	return false;
59	0	}
60
61	0	const std::unique_ptr<rocksdb::DB> db(dbptr);
62
63	0	status = env->NewWritableFile(dump_options.dump_location, &dumpfile,
64	0	rocksdb::EnvOptions());
65	0	if (!status.ok()) {
66	0	std::cerr << "Unable to open dump file '" << dump_options.dump_location
67	0	<< "' for writing: " << status.ToString() << std::endl;
68	0	return false;
69	0	}
70
71	0	rocksdb::Slice magicslice(magicstr, 8);
72	0	status = dumpfile->Append(magicslice);
73	0	if (!status.ok()) {
74	0	std::cerr << "Append failed: " << status.ToString() << std::endl;
75	0	return false;
76	0	}
77
78	0	rocksdb::Slice versionslice(versionstr, 8);
79	0	status = dumpfile->Append(versionslice);
80	0	if (!status.ok()) {
81	0	std::cerr << "Append failed: " << status.ToString() << std::endl;
82	0	return false;
83	0	}
84
85	0	if (dump_options.anonymous) {
86	0	snprintf(json, sizeof(json), "{}");
87	0	} else {
88	0	status = env->GetHostName(hostname, sizeof(hostname));
89	0	status = env->GetCurrentTime(&timesec);
90	0	status = env->GetAbsolutePath(dump_options.db_path, &abspath);
91	0	snprintf(json, sizeof(json),
92	0	"{ \"database-path\": \"%s\", \"hostname\": \"%s\", "
93	0	"\"creation-time\": %" PRIi64 " }",
94	0	abspath.c_str(), hostname, timesec);
95	0	}
96
97	0	rocksdb::Slice infoslice(json, strlen(json));
98	0	char infosize[4];
99	0	rocksdb::EncodeFixed32(infosize, (uint32_t)infoslice.size());
100	0	rocksdb::Slice infosizeslice(infosize, 4);
101	0	status = dumpfile->Append(infosizeslice);
102	0	if (!status.ok()) {
103	0	std::cerr << "Append failed: " << status.ToString() << std::endl;
104	0	return false;
105	0	}
106	0	status = dumpfile->Append(infoslice);
107	0	if (!status.ok()) {
108	0	std::cerr << "Append failed: " << status.ToString() << std::endl;
109	0	return false;
110	0	}
111
112	0	const std::unique_ptr<rocksdb::Iterator> it(
113	0	db->NewIterator(rocksdb::ReadOptions()));
114	0	for (it->SeekToFirst(); it->Valid(); it->Next()) {
115	0	char keysize[4];
116	0	rocksdb::EncodeFixed32(keysize, (uint32_t)it->key().size());
117	0	rocksdb::Slice keysizeslice(keysize, 4);
118	0	status = dumpfile->Append(keysizeslice);
119	0	if (!status.ok()) {
120	0	std::cerr << "Append failed: " << status.ToString() << std::endl;
121	0	return false;
122	0	}
123	0	status = dumpfile->Append(it->key());
124	0	if (!status.ok()) {
125	0	std::cerr << "Append failed: " << status.ToString() << std::endl;
126	0	return false;
127	0	}
128
129	0	char valsize[4];
130	0	rocksdb::EncodeFixed32(valsize, (uint32_t)it->value().size());
131	0	rocksdb::Slice valsizeslice(valsize, 4);
132	0	status = dumpfile->Append(valsizeslice);
133	0	if (!status.ok()) {
134	0	std::cerr << "Append failed: " << status.ToString() << std::endl;
135	0	return false;
136	0	}
137	0	status = dumpfile->Append(it->value());
138	0	if (!status.ok()) {
139	0	std::cerr << "Append failed: " << status.ToString() << std::endl;
140	0	return false;
141	0	}
142	0	}
143	0	if (!it->status().ok()) {
144	0	std::cerr << "Database iteration failed: " << status.ToString()
145	0	<< std::endl;
146	0	return false;
147	0	}
148	0	return true;
149	0	}
150
151		bool DbUndumpTool::Run(const UndumpOptions& undump_options,
152	0	rocksdb::Options options) {
153	0	rocksdb::DB* dbptr;
154	0	rocksdb::Status status;
155	0	rocksdb::Env* env;
156	0	std::unique_ptr<rocksdb::SequentialFile> dumpfile;
157	0	rocksdb::Slice slice;
158	0	uint8_t scratch8[8];
159
160	0	static const char* magicstr = "ROCKDUMP";
161	0	static const char versionstr[8] = {0, 0, 0, 0, 0, 0, 0, 1};
162
163	0	env = rocksdb::Env::Default();
164
165	0	status = env->NewSequentialFile(undump_options.dump_location, &dumpfile,
166	0	rocksdb::EnvOptions());
167	0	if (!status.ok()) {
168	0	std::cerr << "Unable to open dump file '" << undump_options.dump_location
169	0	<< "' for reading: " << status.ToString() << std::endl;
170	0	return false;
171	0	}
172
173	0	status = dumpfile->Read(8, &slice, scratch8);
174	0	if (!status.ok() \|\| slice.size() != 8 \|\|
175	0	memcmp(slice.data(), magicstr, 8) != 0) {
176	0	std::cerr << "File '" << undump_options.dump_location
177	0	<< "' is not a recognizable dump file." << std::endl;
178	0	return false;
179	0	}
180
181	0	status = dumpfile->Read(8, &slice, scratch8);
182	0	if (!status.ok() \|\| slice.size() != 8 \|\|
183	0	memcmp(slice.data(), versionstr, 8) != 0) {
184	0	std::cerr << "File '" << undump_options.dump_location
185	0	<< "' version not recognized." << std::endl;
186	0	return false;
187	0	}
188
189	0	status = dumpfile->Read(4, &slice, scratch8);
190	0	if (!status.ok() \|\| slice.size() != 4) {
191	0	std::cerr << "Unable to read info blob size." << std::endl;
192	0	return false;
193	0	}
194	0	uint32_t infosize = rocksdb::DecodeFixed32(slice.data());
195	0	status = dumpfile->Skip(infosize);
196	0	if (!status.ok()) {
197	0	std::cerr << "Unable to skip info blob: " << status.ToString() << std::endl;
198	0	return false;
199	0	}
200
201	0	options.create_if_missing = true;
202	0	status = rocksdb::DB::Open(options, undump_options.db_path, &dbptr);
203	0	if (!status.ok()) {
204	0	std::cerr << "Unable to open database '" << undump_options.db_path
205	0	<< "' for writing: " << status.ToString() << std::endl;
206	0	return false;
207	0	}
208
209	0	const std::unique_ptr<rocksdb::DB> db(dbptr);
210
211	0	uint32_t last_keysize = 64;
212	0	size_t last_valsize = 1 << 20;
213	0	std::unique_ptr<uint8_t[]> keyscratch(new uint8_t[last_keysize]);
214	0	std::unique_ptr<uint8_t[]> valscratch(new uint8_t[last_valsize]);
215
216	0	while (1) {
217	0	uint32_t keysize, valsize;
218	0	rocksdb::Slice keyslice;
219	0	rocksdb::Slice valslice;
220
221	0	status = dumpfile->Read(4, &slice, scratch8);
222	0	if (!status.ok() \|\| slice.size() != 4) break;
223	0	keysize = rocksdb::DecodeFixed32(slice.data());
224	0	if (keysize > last_keysize) {
225	0	while (keysize > last_keysize) last_keysize *= 2;
226	0	keyscratch = std::unique_ptr<uint8_t[]>(new uint8_t[last_keysize]);
227	0	}
228
229	0	status = dumpfile->Read(keysize, &keyslice, keyscratch.get());
230	0	if (!status.ok() \|\| keyslice.size() != keysize) {
231	0	std::cerr << "Key read failure: "
232	0	<< (status.ok() ? "insufficient data" : status.ToString())
233	0	<< std::endl;
234	0	return false;
235	0	}
236
237	0	status = dumpfile->Read(4, &slice, scratch8);
238	0	if (!status.ok() \|\| slice.size() != 4) {
239	0	std::cerr << "Unable to read value size: "
240	0	<< (status.ok() ? "insufficient data" : status.ToString())
241	0	<< std::endl;
242	0	return false;
243	0	}
244	0	valsize = rocksdb::DecodeFixed32(slice.data());
245	0	if (valsize > last_valsize) {
246	0	while (valsize > last_valsize) last_valsize *= 2;
247	0	valscratch = std::unique_ptr<uint8_t[]>(new uint8_t[last_valsize]);
248	0	}
249
250	0	status = dumpfile->Read(valsize, &valslice, valscratch.get());
251	0	if (!status.ok() \|\| valslice.size() != valsize) {
252	0	std::cerr << "Unable to read value: "
253	0	<< (status.ok() ? "insufficient data" : status.ToString())
254	0	<< std::endl;
255	0	return false;
256	0	}
257
258	0	status = db->Put(rocksdb::WriteOptions(), keyslice, valslice);
259	0	if (!status.ok()) {
260	0	fprintf(stderr, "Unable to write database entry\n");
261	0	return false;
262	0	}
263	0	}
264
265	0	if (undump_options.compact_db) {
266	0	status = db->CompactRange(rocksdb::CompactRangeOptions(), nullptr, nullptr);
267	0	if (!status.ok()) {
268	0	fprintf(stderr,
269	0	"Unable to compact the database after loading the dumped file\n");
270	0	return false;
271	0	}
272	0	}
273	0	return true;
274	0	}
275		} // namespace rocksdb
276		#endif // ROCKSDB_LITE