/Users/deen/code/yugabyte-db/src/yb/rocksdb/tools/dump/db_dump_tool.cc
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(×ec); |
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 |