#include <rocksdb/db.h>
|
#include <rocksdb/options.h>
|
#include <iostream>
|
#include <string>
|
#include <unordered_map>
|
#include <vector>
|
#include <regex>
|
#include "tser/tser.hpp"
|
|
// Wrapper class for RocksDB
|
class RocksDBWrapper
|
{
|
public:
|
RocksDBWrapper(const std::string &db_path, const std::string &index_path);
|
|
~RocksDBWrapper();
|
|
template <typename T>
|
void store(const std::string &key, const T &value)
|
{
|
tser::BinaryArchive archive(0);
|
archive.save(value);
|
|
rocksdb::Status status = db_->Put(rocksdb::WriteOptions(), serializeKey(key), archive.get_buffer().data());
|
if (!status.ok())
|
{
|
std::cerr << "Error storing value in RocksDB database: " << status.ToString() << std::endl;
|
}
|
|
// Index the struct's members
|
index_members(serializeKey(key), value);
|
}
|
|
template <typename T>
|
bool get(const std::string &key, T &value)
|
{
|
std::string serialized_value;
|
rocksdb::Status status = db_->Get(rocksdb::ReadOptions(), serializeKey(key), &serialized_value);
|
if (!status.ok())
|
{
|
std::cerr << "Error getting value from RocksDB database: " << status.ToString() << std::endl;
|
return false;
|
}
|
|
tser::BinaryArchive archive(0);
|
archive.initialize(serialized_value);
|
value = archive.load<T>();
|
return true;
|
}
|
|
template <typename T>
|
std::vector<std::string> search(const std::string &member_name, const std::string &member_value)
|
{
|
std::vector<std::string> keys;
|
std::string prefix = member_name + ":" + serializeKey(member_value) + ":";
|
rocksdb::Iterator *it = index_db_->NewIterator(rocksdb::ReadOptions());
|
for (it->Seek(prefix); it->Valid() && it->key().ToString().find(prefix) == 0; it->Next())
|
{
|
keys.push_back(it->value().ToString());
|
}
|
delete it;
|
return keys;
|
}
|
|
inline std::vector<std::string> split(const std::string &s, const std::string &delimiter)
|
{
|
size_t last = 0;
|
size_t next = 0;
|
std::vector<std::string> results;
|
while ((next = s.find(delimiter, last)) != std::string::npos)
|
{
|
results.emplace_back(s.substr(last, next - last));
|
last = next + delimiter.length();
|
}
|
results.emplace_back(s.substr(last));
|
return results;
|
}
|
|
template <typename T>
|
void search_text(const std::string &search_value, std::vector<T> &res)
|
{
|
std::vector<std::string> keys;
|
std::vector<std::string> tmp;
|
rocksdb::Iterator *it = index_db_->NewIterator(rocksdb::ReadOptions());
|
|
for (it->SeekToFirst(); it->Valid(); it->Next())
|
{
|
std::string key = it->key().ToString();
|
auto _split = this->split(key, ":");
|
if (this->StringMatch(_split[1], search_value))
|
{
|
keys.push_back(_split[2]);
|
}
|
}
|
|
std::vector<rocksdb::Slice> slices;
|
for (const auto &key : keys)
|
{
|
slices.push_back(rocksdb::Slice(key));
|
}
|
|
std::vector<rocksdb::Status> statuses = db_->MultiGet(rocksdb::ReadOptions(), slices, &tmp);
|
|
for (const auto &_v : tmp)
|
{
|
T _val;
|
this->DeSerialize(_v, _val);
|
res.emplace_back(_val);
|
}
|
delete it;
|
}
|
|
// New keyExists method
|
bool keyExists(const std::string &key);
|
|
// New conditional search method
|
template <typename T>
|
void search_conditional(std::vector<T> &res, std::function<bool(const T &)> condition)
|
{
|
rocksdb::Iterator *it = db_->NewIterator(rocksdb::ReadOptions());
|
|
for (it->SeekToFirst(); it->Valid(); it->Next())
|
{
|
std::string key = it->key().ToString();
|
std::string serialized_value = it->value().ToString();
|
|
T _val;
|
this->DeSerialize(serialized_value, _val);
|
|
if (condition(_val))
|
{
|
res.push_back(_val);
|
}
|
}
|
delete it;
|
}
|
|
private:
|
rocksdb::DB *db_;
|
rocksdb::DB *index_db_;
|
|
bool StringMatch(const std::string &text, const std::string &pattern);
|
|
template <typename T>
|
inline void DeSerialize(const std::string &data, T &value)
|
{
|
tser::BinaryArchive archive(0);
|
archive.initialize(data);
|
value = archive.load<T>();
|
}
|
|
template <typename T>
|
inline void Serialize(const T &value, std::string &data)
|
{
|
tser::BinaryArchive archive(0);
|
archive.save(value);
|
data = archive.get_buffer().data();
|
}
|
|
template <typename T>
|
void index_single_member(const std::string &key, const std::string &member_name, const T &member_value)
|
{
|
if constexpr (std::is_arithmetic_v<std::decay_t<T>>)
|
{
|
std::string member_value_str = std::to_string(member_value);
|
std::string index_key = member_name + ":" + serializeKey(member_value_str) + ":" + key;
|
rocksdb::Status status = index_db_->Put(rocksdb::WriteOptions(), index_key, key);
|
if (!status.ok())
|
{
|
std::cerr << "Error indexing member in RocksDB index database: " << status.ToString() << std::endl;
|
}
|
}
|
else if constexpr (std::is_same_v<std::decay_t<T>, std::string>)
|
{
|
std::string index_key = member_name + ":" + serializeKey(member_value) + ":" + key;
|
rocksdb::Status status = index_db_->Put(rocksdb::WriteOptions(), index_key, key);
|
if (!status.ok())
|
{
|
std::cerr << "Error indexing member in RocksDB index database: " << status.ToString() << std::endl;
|
}
|
}
|
else if constexpr (tser::is_container_v<std::decay_t<T>>)
|
{
|
for (const auto &elem : member_value)
|
{
|
index_single_member(key, member_name, elem);
|
}
|
}
|
else
|
{
|
// Handle other types if necessary
|
}
|
}
|
|
template <typename T>
|
void index_members(const std::string &key, const T &value)
|
{
|
if constexpr (tser::is_tser_t_v<T>)
|
{
|
auto member_names = T::_memberNames;
|
size_t index = 0;
|
std::apply([&](auto &&...args)
|
{ ((index_single_member(key, member_names[index++], args)), ...); },
|
value.members());
|
}
|
}
|
|
std::string serializeKey(const std::string &key)
|
{
|
// Implement a method to serialize keys to handle special characters
|
std::string serialized_key;
|
for (const char c : key)
|
{
|
if (c == ':')
|
{
|
serialized_key += "%3A"; // Percent-encoding for ':'
|
}
|
else
|
{
|
serialized_key += c;
|
}
|
}
|
return serialized_key;
|
}
|
};
|