#include #include #include #include #include #include "db_wrapper.h" #include "rocksdb/db.h" #include "rocksdb/slice.h" #include "rocksdb/options.h" namespace { void printWithBackSlashes(std::string str) { for (std::string::size_type i = 0; i < str.size(); i++) { if (str[i] == '\\' || str[i] == '"') { std::cout << "\\"; } std::cout << str[i]; } } bool has_key_for_array(Local obj, std::string key) { return obj->Has(String::NewSymbol(key.c_str())) && obj->Get(String::NewSymbol(key.c_str()))->IsArray(); } } using namespace v8; Persistent DBWrapper::constructor; DBWrapper::DBWrapper() { options_.IncreaseParallelism(); options_.OptimizeLevelStyleCompaction(); options_.disable_auto_compactions = true; options_.create_if_missing = true; } DBWrapper::~DBWrapper() { delete db_; } bool DBWrapper::HasFamilyNamed(std::string& name, DBWrapper* db) { return db->columnFamilies_.find(name) != db->columnFamilies_.end(); } void DBWrapper::Init(Handle exports) { Local tpl = FunctionTemplate::New(New); tpl->SetClassName(String::NewSymbol("DBWrapper")); tpl->InstanceTemplate()->SetInternalFieldCount(8); tpl->PrototypeTemplate()->Set(String::NewSymbol("open"), FunctionTemplate::New(Open)->GetFunction()); tpl->PrototypeTemplate()->Set(String::NewSymbol("get"), FunctionTemplate::New(Get)->GetFunction()); tpl->PrototypeTemplate()->Set(String::NewSymbol("put"), FunctionTemplate::New(Put)->GetFunction()); tpl->PrototypeTemplate()->Set(String::NewSymbol("delete"), FunctionTemplate::New(Delete)->GetFunction()); tpl->PrototypeTemplate()->Set(String::NewSymbol("dump"), FunctionTemplate::New(Dump)->GetFunction()); tpl->PrototypeTemplate()->Set(String::NewSymbol("createColumnFamily"), FunctionTemplate::New(CreateColumnFamily)->GetFunction()); tpl->PrototypeTemplate()->Set(String::NewSymbol("writeBatch"), FunctionTemplate::New(WriteBatch)->GetFunction()); tpl->PrototypeTemplate()->Set(String::NewSymbol("compactRange"), FunctionTemplate::New(CompactRange)->GetFunction()); constructor = Persistent::New(tpl->GetFunction()); exports->Set(String::NewSymbol("DBWrapper"), constructor); } Handle DBWrapper::Open(const Arguments& args) { HandleScope scope; DBWrapper* db_wrapper = ObjectWrap::Unwrap(args.This()); if (!(args[0]->IsString() && (args[1]->IsUndefined() || args[1]->IsArray()))) { return scope.Close(Boolean::New(false)); } std::string db_file = *v8::String::Utf8Value(args[0]->ToString()); std::vector cfs = { rocksdb::kDefaultColumnFamilyName }; if (!args[1]->IsUndefined()) { Handle array = Handle::Cast(args[1]); for (uint i = 0; i < array->Length(); i++) { if (!array->Get(i)->IsString()) { return scope.Close(Boolean::New(false)); } cfs.push_back(*v8::String::Utf8Value(array->Get(i)->ToString())); } } if (cfs.size() == 1) { db_wrapper->status_ = rocksdb::DB::Open( db_wrapper->options_, db_file, &db_wrapper->db_); return scope.Close(Boolean::New(db_wrapper->status_.ok())); } std::vector families; for (std::vector::size_type i = 0; i < cfs.size(); i++) { families.push_back(rocksdb::ColumnFamilyDescriptor( cfs[i], rocksdb::ColumnFamilyOptions())); } std::vector handles; db_wrapper->status_ = rocksdb::DB::Open( db_wrapper->options_, db_file, families, &handles, &db_wrapper->db_); if (!db_wrapper->status_.ok()) { return scope.Close(Boolean::New(db_wrapper->status_.ok())); } for (std::vector::size_type i = 0; i < handles.size(); i++) { db_wrapper->columnFamilies_[cfs[i]] = handles[i]; } return scope.Close(Boolean::New(true)); } Handle DBWrapper::New(const Arguments& args) { HandleScope scope; Handle to_return; if (args.IsConstructCall()) { DBWrapper* db_wrapper = new DBWrapper(); db_wrapper->Wrap(args.This()); return args.This(); } const int argc = 0; Local argv[0] = {}; return scope.Close(constructor->NewInstance(argc, argv)); } Handle DBWrapper::Get(const Arguments& args) { HandleScope scope; if (!(args[0]->IsString() && (args[1]->IsUndefined() || args[1]->IsString()))) { return scope.Close(Null()); } DBWrapper* db_wrapper = ObjectWrap::Unwrap(args.This()); std::string key = *v8::String::Utf8Value(args[0]->ToString()); std::string cf = *v8::String::Utf8Value(args[1]->ToString()); std::string value; if (args[1]->IsUndefined()) { db_wrapper->status_ = db_wrapper->db_->Get( rocksdb::ReadOptions(), key, &value); } else if (db_wrapper->HasFamilyNamed(cf, db_wrapper)) { db_wrapper->status_ = db_wrapper->db_->Get( rocksdb::ReadOptions(), db_wrapper->columnFamilies_[cf], key, &value); } else { return scope.Close(Null()); } Handle v = db_wrapper->status_.ok() ? String::NewSymbol(value.c_str()) : Null(); return scope.Close(v); } Handle DBWrapper::Put(const Arguments& args) { HandleScope scope; if (!(args[0]->IsString() && args[1]->IsString() && (args[2]->IsUndefined() || args[2]->IsString()))) { return scope.Close(Boolean::New(false)); } DBWrapper* db_wrapper = ObjectWrap::Unwrap(args.This()); std::string key = *v8::String::Utf8Value(args[0]->ToString()); std::string value = *v8::String::Utf8Value(args[1]->ToString()); std::string cf = *v8::String::Utf8Value(args[2]->ToString()); if (args[2]->IsUndefined()) { db_wrapper->status_ = db_wrapper->db_->Put( rocksdb::WriteOptions(), key, value ); } else if (db_wrapper->HasFamilyNamed(cf, db_wrapper)) { db_wrapper->status_ = db_wrapper->db_->Put( rocksdb::WriteOptions(), db_wrapper->columnFamilies_[cf], key, value ); } else { return scope.Close(Boolean::New(false)); } return scope.Close(Boolean::New(db_wrapper->status_.ok())); } Handle DBWrapper::Delete(const Arguments& args) { HandleScope scope; if (!args[0]->IsString()) { return scope.Close(Boolean::New(false)); } DBWrapper* db_wrapper = ObjectWrap::Unwrap(args.This()); std::string arg0 = *v8::String::Utf8Value(args[0]->ToString()); std::string arg1 = *v8::String::Utf8Value(args[1]->ToString()); if (args[1]->IsUndefined()) { db_wrapper->status_ = db_wrapper->db_->Delete( rocksdb::WriteOptions(), arg0); } else { if (!db_wrapper->HasFamilyNamed(arg1, db_wrapper)) { return scope.Close(Boolean::New(false)); } db_wrapper->status_ = db_wrapper->db_->Delete( rocksdb::WriteOptions(), db_wrapper->columnFamilies_[arg1], arg0); } return scope.Close(Boolean::New(db_wrapper->status_.ok())); } Handle DBWrapper::Dump(const Arguments& args) { HandleScope scope; std::unique_ptr iterator; DBWrapper* db_wrapper = ObjectWrap::Unwrap(args.This()); std::string arg0 = *v8::String::Utf8Value(args[0]->ToString()); if (args[0]->IsUndefined()) { iterator.reset(db_wrapper->db_->NewIterator(rocksdb::ReadOptions())); } else { if (!db_wrapper->HasFamilyNamed(arg0, db_wrapper)) { return scope.Close(Boolean::New(false)); } iterator.reset(db_wrapper->db_->NewIterator( rocksdb::ReadOptions(), db_wrapper->columnFamilies_[arg0])); } iterator->SeekToFirst(); while (iterator->Valid()) { std::cout << "\""; printWithBackSlashes(iterator->key().ToString()); std::cout << "\" => \""; printWithBackSlashes(iterator->value().ToString()); std::cout << "\"\n"; iterator->Next(); } return scope.Close(Boolean::New(true)); } Handle DBWrapper::CreateColumnFamily(const Arguments& args) { HandleScope scope; if (!args[0]->IsString()) { return scope.Close(Boolean::New(false)); } DBWrapper* db_wrapper = ObjectWrap::Unwrap(args.This()); std::string cf_name = *v8::String::Utf8Value(args[0]->ToString()); if (db_wrapper->HasFamilyNamed(cf_name, db_wrapper)) { return scope.Close(Boolean::New(false)); } rocksdb::ColumnFamilyHandle* cf; db_wrapper->status_ = db_wrapper->db_->CreateColumnFamily( rocksdb::ColumnFamilyOptions(), cf_name, &cf); if (!db_wrapper->status_.ok()) { return scope.Close(Boolean::New(false)); } db_wrapper->columnFamilies_[cf_name] = cf; return scope.Close(Boolean::New(true)); } bool DBWrapper::AddToBatch(rocksdb::WriteBatch& batch, bool del, Handle array) { Handle put_pair; for (uint i = 0; i < array->Length(); i++) { if (del) { if (!array->Get(i)->IsString()) { return false; } batch.Delete(*v8::String::Utf8Value(array->Get(i)->ToString())); continue; } if (!array->Get(i)->IsArray()) { return false; } put_pair = Handle::Cast(array->Get(i)); if (!put_pair->Get(0)->IsString() || !put_pair->Get(1)->IsString()) { return false; } batch.Put( *v8::String::Utf8Value(put_pair->Get(0)->ToString()), *v8::String::Utf8Value(put_pair->Get(1)->ToString())); } return true; } bool DBWrapper::AddToBatch(rocksdb::WriteBatch& batch, bool del, Handle array, DBWrapper* db_wrapper, std::string cf) { Handle put_pair; for (uint i = 0; i < array->Length(); i++) { if (del) { if (!array->Get(i)->IsString()) { return false; } batch.Delete( db_wrapper->columnFamilies_[cf], *v8::String::Utf8Value(array->Get(i)->ToString())); continue; } if (!array->Get(i)->IsArray()) { return false; } put_pair = Handle::Cast(array->Get(i)); if (!put_pair->Get(0)->IsString() || !put_pair->Get(1)->IsString()) { return false; } batch.Put( db_wrapper->columnFamilies_[cf], *v8::String::Utf8Value(put_pair->Get(0)->ToString()), *v8::String::Utf8Value(put_pair->Get(1)->ToString())); } return true; } Handle DBWrapper::WriteBatch(const Arguments& args) { HandleScope scope; if (!args[0]->IsArray()) { return scope.Close(Boolean::New(false)); } DBWrapper* db_wrapper = ObjectWrap::Unwrap(args.This()); Handle sub_batches = Handle::Cast(args[0]); Local sub_batch; rocksdb::WriteBatch batch; bool well_formed; for (uint i = 0; i < sub_batches->Length(); i++) { if (!sub_batches->Get(i)->IsObject()) { return scope.Close(Boolean::New(false)); } sub_batch = sub_batches->Get(i)->ToObject(); if (sub_batch->Has(String::NewSymbol("column_family"))) { if (!has_key_for_array(sub_batch, "put") && !has_key_for_array(sub_batch, "delete")) { return scope.Close(Boolean::New(false)); } well_formed = db_wrapper->AddToBatch( batch, false, Handle::Cast(sub_batch->Get(String::NewSymbol("put"))), db_wrapper, *v8::String::Utf8Value(sub_batch->Get( String::NewSymbol("column_family")))); well_formed = db_wrapper->AddToBatch( batch, true, Handle::Cast(sub_batch->Get(String::NewSymbol("delete"))), db_wrapper, *v8::String::Utf8Value(sub_batch->Get( String::NewSymbol("column_family")))); } else { well_formed = db_wrapper->AddToBatch( batch, false, Handle::Cast(sub_batch->Get(String::NewSymbol("put")))); well_formed = db_wrapper->AddToBatch( batch, true, Handle::Cast(sub_batch->Get(String::NewSymbol("delete")))); if (!well_formed) { return scope.Close(Boolean::New(false)); } } } db_wrapper->status_ = db_wrapper->db_->Write(rocksdb::WriteOptions(), &batch); return scope.Close(Boolean::New(db_wrapper->status_.ok())); } Handle DBWrapper::CompactRangeDefault(const Arguments& args) { HandleScope scope; DBWrapper* db_wrapper = ObjectWrap::Unwrap(args.This()); rocksdb::Slice begin = *v8::String::Utf8Value(args[0]->ToString()); rocksdb::Slice end = *v8::String::Utf8Value(args[1]->ToString()); db_wrapper->status_ = db_wrapper->db_->CompactRange(&end, &begin); return scope.Close(Boolean::New(db_wrapper->status_.ok())); } Handle DBWrapper::CompactColumnFamily(const Arguments& args) { HandleScope scope; DBWrapper* db_wrapper = ObjectWrap::Unwrap(args.This()); rocksdb::Slice begin = *v8::String::Utf8Value(args[0]->ToString()); rocksdb::Slice end = *v8::String::Utf8Value(args[1]->ToString()); std::string cf = *v8::String::Utf8Value(args[2]->ToString()); db_wrapper->status_ = db_wrapper->db_->CompactRange( db_wrapper->columnFamilies_[cf], &begin, &end); return scope.Close(Boolean::New(db_wrapper->status_.ok())); } Handle DBWrapper::CompactOptions(const Arguments& args) { HandleScope scope; if (!args[2]->IsObject()) { return scope.Close(Boolean::New(false)); } DBWrapper* db_wrapper = ObjectWrap::Unwrap(args.This()); rocksdb::Slice begin = *v8::String::Utf8Value(args[0]->ToString()); rocksdb::Slice end = *v8::String::Utf8Value(args[1]->ToString()); Local options = args[2]->ToObject(); int target_level = -1, target_path_id = 0; if (options->Has(String::NewSymbol("target_level")) && options->Get(String::NewSymbol("target_level"))->IsInt32()) { target_level = (int)(options->Get( String::NewSymbol("target_level"))->ToInt32()->Value()); if (options->Has(String::NewSymbol("target_path_id")) || options->Get(String::NewSymbol("target_path_id"))->IsInt32()) { target_path_id = (int)(options->Get( String::NewSymbol("target_path_id"))->ToInt32()->Value()); } } db_wrapper->status_ = db_wrapper->db_->CompactRange( &begin, &end, true, target_level, target_path_id ); return scope.Close(Boolean::New(db_wrapper->status_.ok())); } Handle DBWrapper::CompactAll(const Arguments& args) { HandleScope scope; if (!args[2]->IsObject() || !args[3]->IsString()) { return scope.Close(Boolean::New(false)); } DBWrapper* db_wrapper = ObjectWrap::Unwrap(args.This()); rocksdb::Slice begin = *v8::String::Utf8Value(args[0]->ToString()); rocksdb::Slice end = *v8::String::Utf8Value(args[1]->ToString()); Local options = args[2]->ToObject(); std::string cf = *v8::String::Utf8Value(args[3]->ToString()); int target_level = -1, target_path_id = 0; if (options->Has(String::NewSymbol("target_level")) && options->Get(String::NewSymbol("target_level"))->IsInt32()) { target_level = (int)(options->Get( String::NewSymbol("target_level"))->ToInt32()->Value()); if (options->Has(String::NewSymbol("target_path_id")) || options->Get(String::NewSymbol("target_path_id"))->IsInt32()) { target_path_id = (int)(options->Get( String::NewSymbol("target_path_id"))->ToInt32()->Value()); } } db_wrapper->status_ = db_wrapper->db_->CompactRange( db_wrapper->columnFamilies_[cf], &begin, &end, true, target_level, target_path_id); return scope.Close(Boolean::New(db_wrapper->status_.ok())); } Handle DBWrapper::CompactRange(const Arguments& args) { HandleScope scope; if (!args[0]->IsString() || !args[1]->IsString()) { return scope.Close(Boolean::New(false)); } switch(args.Length()) { case 2: return CompactRangeDefault(args); case 3: return args[2]->IsString() ? CompactColumnFamily(args) : CompactOptions(args); default: return CompactAll(args); } } Handle DBWrapper::Close(const Arguments& args) { HandleScope scope; delete ObjectWrap::Unwrap(args.This()); return scope.Close(Null()); }