You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

235 lines
6.3KB

  1. /** zstreedb2 main.cxx
  2. (C) 2018 Erik Zscheile
  3. License: GPLv2
  4. **/
  5. #include <stdio.h>
  6. #include <stdlib.h>
  7. #include <algorithm>
  8. #include <thread>
  9. #include <unordered_map>
  10. #include <utility>
  11. #include <rocksdb/db.h>
  12. #include "args.hpp"
  13. #include "error.hpp"
  14. using namespace std;
  15. const char *progname;
  16. static rocksdb::DB* db;
  17. static void cleanup() {
  18. delete db;
  19. }
  20. static void filter_match(vector<string> &selv, string &&key, const string &value) {
  21. vector<rocksdb::Status> sts;
  22. vector<string> values;
  23. // 1. retrieve values
  24. {
  25. string sufx = "." + move(key);
  26. vector<string> keys;
  27. vector<rocksdb::Slice> key_slices;
  28. keys.reserve(selv.size());
  29. key_slices.reserve(selv.size());
  30. for(const auto &i : selv)
  31. key_slices.emplace_back(keys.emplace_back(i + sufx));
  32. string().swap(sufx);
  33. sts = db->MultiGet({}, key_slices, &values);
  34. }
  35. // 2. find matches
  36. {
  37. auto it_s = sts.begin();
  38. auto it_v = values.begin();
  39. auto it_sel = selv.begin();
  40. while(it_sel != selv.end()) {
  41. if(!it_s->ok() || *it_v != value) {
  42. it_s = sts.erase(it_s);
  43. it_v = values.erase(it_v);
  44. it_sel = selv.erase(it_sel);
  45. continue;
  46. }
  47. ++it_s; ++it_v; ++it_sel;
  48. }
  49. }
  50. }
  51. #define AC(C) case ztdb::args_t::C:
  52. int main(const int argc, char *argv[]) {
  53. progname = argv[0];
  54. ztdb::args_t args;
  55. if(!args.parse(argc, argv)) {
  56. ztdb::help();
  57. errmsg(1, "invalid arguments passed");
  58. }
  59. rocksdb::Status st;
  60. {
  61. rocksdb::Options opts;
  62. switch(args.cmd) {
  63. AC(MV) AC(RM)
  64. opts.IncreaseParallelism(std::max(thread::hardware_concurrency(), 1u));
  65. default: break;
  66. }
  67. opts.allow_mmap_reads = true;
  68. opts.allow_mmap_writes = true;
  69. opts.compression = args.compress;
  70. opts.bottommost_compression = args.compress;
  71. opts.create_if_missing = true;
  72. st = args.flag_use_ro()
  73. ? rocksdb::DB::OpenForReadOnly(opts, ".", &db)
  74. : rocksdb::DB::Open (opts, ".", &db);
  75. if(!st.ok())
  76. errmsg(2, "db open failed: " + st.ToString());
  77. }
  78. atexit(cleanup);
  79. switch(args.cmd) {
  80. AC(GET)
  81. {
  82. rocksdb::PinnableSlice slc;
  83. st = db->Get({}, db->DefaultColumnFamily(), args.sel_node, &slc);
  84. if(st.ok())
  85. printf("%.*s\n", static_cast<int>(slc.size()), slc.data());
  86. }
  87. break;
  88. AC(SET)
  89. if(args.cmd_args.empty()) {
  90. db->Delete({}, args.sel_node);
  91. } else {
  92. st = db->Put({}, args.sel_node, args.cmd_args.front());
  93. if(!st.ok()) errmsg(2, "db put failed: " + st.ToString());
  94. }
  95. break;
  96. AC(LS)
  97. {
  98. const bool sall = (args.sel_node == "*");
  99. const string xpar = args.sel_node + '.';
  100. const auto it = db->NewIterator({});
  101. if(sall) it->SeekToFirst();
  102. else it->Seek(xpar);
  103. for(; it->Valid(); it->Next()) {
  104. auto tmp = it->key();
  105. if(!sall) {
  106. if(!tmp.starts_with(xpar)) break;
  107. tmp.remove_prefix(xpar.size());
  108. }
  109. printf("%.*s\n", static_cast<int>(tmp.size()), tmp.data());
  110. }
  111. delete it;
  112. }
  113. break;
  114. AC(PRINT)
  115. {
  116. if(args.cmd_args.size() < 3)
  117. errmsg(1, "invalid invocation to 'zstreedb2 p PREFIX BFK BFV AFV'");
  118. const bool sall = (args.sel_node == "*");
  119. const string xpar = args.sel_node + '.';
  120. const char *bfk = args.cmd_args[0].c_str();
  121. const char *bfv = args.cmd_args[1].c_str();
  122. const char *afv = args.cmd_args[2].c_str();
  123. const auto it = db->NewIterator({});
  124. if(sall) it->SeekToFirst();
  125. else it->Seek(xpar);
  126. for(; it->Valid(); it->Next()) {
  127. auto tmp = it->key();
  128. if(!sall) {
  129. if(!tmp.starts_with(xpar)) break;
  130. tmp.remove_prefix(xpar.size());
  131. }
  132. auto val = it->value();
  133. printf("%s%.*s%s%.*s%s\n", bfk, static_cast<int>(tmp.size()), tmp.data(),
  134. bfv, static_cast<int>(val.size()), val.data(), afv);
  135. }
  136. delete it;
  137. }
  138. break;
  139. AC(MV)
  140. if(!args.cmd_args.empty() && args.sel_node != args.cmd_args.front()) {
  141. const auto it = db->NewIterator({});
  142. const string xpar = args.sel_node + '.';
  143. rocksdb::WriteBatch batch;
  144. for(it->Seek(xpar); it->Valid() && it->key().starts_with(xpar); it->Next()) {
  145. batch.Delete(it->key());
  146. string tmp = it->key().ToString();
  147. tmp.replace(0, args.sel_node.size(), args.cmd_args.front());
  148. batch.Put(tmp, it->value());
  149. }
  150. st = db->Write({}, &batch);
  151. delete it;
  152. if(!st.ok()) errmsg(2, "db write failed: " + st.ToString());
  153. } else {
  154. errmsg(1, "invalid invocation to 'zstreedb2 v FROM TO'");
  155. }
  156. break;
  157. AC(RM)
  158. {
  159. const auto it = db->NewIterator({});
  160. const string xpar = args.sel_node + ".";
  161. rocksdb::WriteBatch batch;
  162. // we could use DeleteRange, but it would be more complex,
  163. // because we would need to handle invalid it after loop (EOF reached)
  164. for(it->Seek(xpar); it->Valid() && it->key().starts_with(xpar); it->Next())
  165. batch.Delete(it->key());
  166. batch.Delete(args.sel_node);
  167. st = db->Write({}, &batch);
  168. delete it;
  169. if(!st.ok()) errmsg(2, "db write failed: " + st.ToString());
  170. }
  171. break;
  172. AC(MATCH)
  173. {
  174. if(args.cmd_args.empty())
  175. errmsg(1, "invalid invocation to 'zstreedb2 m (KEY VALUE)+'");
  176. vector<string> selv;
  177. {
  178. const string key = "." + args.sel_node;
  179. const auto it = db->NewIterator({});
  180. for(it->SeekToFirst(); it->Valid(); it->Next())
  181. if(it->key().ends_with(key) && it->value() == args.cmd_args.front()) {
  182. auto slc = it->key();
  183. slc.remove_suffix(key.size());
  184. selv.emplace_back(slc.data(), slc.size());
  185. }
  186. delete it;
  187. }
  188. args.cmd_args.erase(args.cmd_args.begin());
  189. auto it = args.cmd_args.begin();
  190. while(it != args.cmd_args.end() && (it + 1) != args.cmd_args.end()) {
  191. filter_match(selv, std::move(*it), *(it + 1));
  192. it += 2;
  193. }
  194. for(const auto &i : selv)
  195. printf("%.*s\n", static_cast<int>(i.size()), i.data());
  196. }
  197. break;
  198. AC(UNKNOWN)
  199. errmsg(1, "unknown command");
  200. }
  201. return 0;
  202. }