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