Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: kv: remove key and list prefix #1994

Merged
merged 3 commits into from
Dec 10, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 18 additions & 3 deletions src/kv/kv.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ void KVStore::set(const std::string &key, const ValueWrapper &value) {
}

void KVStore::setLocked(const std::string &key, const ValueWrapper &value) {
logger.debug("KVStore::set({})", key);
logger.trace("KVStore::set({})", key);
auto it = store_.find(key);
if (it != store_.end()) {
it->second.first = value;
Expand All @@ -53,7 +53,7 @@ void KVStore::setLocked(const std::string &key, const ValueWrapper &value) {
}

std::optional<ValueWrapper> KVStore::get(const std::string &key, bool forceLoad /*= false */) {
logger.debug("KVStore::get({})", key);
logger.trace("KVStore::get({})", key);
std::scoped_lock lock(mutex_);
if (forceLoad || !store_.contains(key)) {
auto value = load(key);
Expand All @@ -72,11 +72,26 @@ std::optional<ValueWrapper> KVStore::get(const std::string &key, bool forceLoad
return value;
}

std::unordered_set<std::string> KVStore::keys(const std::string &prefix /*= ""*/) {
std::scoped_lock lock(mutex_);
std::unordered_set<std::string> keys;
for (const auto &[key, value] : store_) {
if (key.find(prefix) == 0) {
std::string suffix = key.substr(prefix.size());
keys.insert(suffix);
}
}
for (const auto &key : loadPrefix(prefix)) {
keys.insert(key);
}
return keys;
}

void KV::remove(const std::string &key) {
set(key, ValueWrapper::deleted());
}

std::shared_ptr<KV> KVStore::scoped(const std::string &scope) {
logger.debug("KVStore::scoped({})", scope);
logger.trace("KVStore::scoped({})", scope);
return std::make_shared<ScopedKV>(logger, *this, scope);
}
10 changes: 9 additions & 1 deletion src/kv/kv.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,8 @@ class KV : public std::enable_shared_from_this<KV> {

virtual std::shared_ptr<KV> scoped(const std::string &scope) = 0;

virtual std::unordered_set<std::string> keys(const std::string &prefix = "") = 0;

void remove(const std::string &key);

virtual void flush() {
Expand Down Expand Up @@ -60,6 +62,7 @@ class KVStore : public KV {
}

std::shared_ptr<KV> scoped(const std::string &scope) override final;
std::unordered_set<std::string> keys(const std::string &prefix = "");

protected:
phmap::parallel_flat_hash_map<std::string, std::pair<ValueWrapper, std::list<std::string>::iterator>> getStore() {
Expand All @@ -76,6 +79,7 @@ class KVStore : public KV {

virtual std::optional<ValueWrapper> load(const std::string &key) = 0;
virtual bool save(const std::string &key, const ValueWrapper &value) = 0;
virtual std::vector<std::string> loadPrefix(const std::string &prefix = "") = 0;

private:
void setLocked(const std::string &key, const ValueWrapper &value);
Expand Down Expand Up @@ -118,10 +122,14 @@ class ScopedKV final : public KV {
}

std::shared_ptr<KV> scoped(const std::string &scope) override final {
logger.debug("ScopedKV::scoped({})", buildKey(scope));
logger.trace("ScopedKV::scoped({})", buildKey(scope));
return std::make_shared<ScopedKV>(logger, rootKV_, buildKey(scope));
}

std::unordered_set<std::string> keys(const std::string &prefix = "") override {
return rootKV_.keys(buildKey(prefix));
}

private:
std::string buildKey(const std::string &key) const {
return fmt::format("{}.{}", prefix_, key);
Expand Down
20 changes: 19 additions & 1 deletion src/kv/kv_sql.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,24 @@ std::optional<ValueWrapper> KVSQL::load(const std::string &key) {
return std::nullopt;
}

std::vector<std::string> KVSQL::loadPrefix(const std::string &prefix /* = ""*/) {
std::vector<std::string> keys;
std::string keySearch = db.escapeString(prefix + "%");
auto query = fmt::format("SELECT `key_name` FROM `kv_store` WHERE `key_name` LIKE {}", keySearch);
auto result = db.storeQuery(query);
if (result == nullptr) {
return keys;
}

do {
std::string key = result->getString("key_name");
replaceString(key, prefix, "");
keys.push_back(key);
} while (result->next());

return keys;
}

bool KVSQL::save(const std::string &key, const ValueWrapper &value) {
auto update = dbUpdate();
prepareSave(key, value, update);
Expand All @@ -58,7 +76,7 @@ bool KVSQL::prepareSave(const std::string &key, const ValueWrapper &value, DBIns
return db.executeQuery(query);
}

update.addRow(fmt::format("{}, {}, {}", db.escapeString(key), getTimeMsNow(), db.escapeString(data)));
update.addRow(fmt::format("{}, {}, {}", db.escapeString(key), value.getTimestamp(), db.escapeString(data)));
return true;
}

Expand Down
1 change: 1 addition & 0 deletions src/kv/kv_sql.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ class KVSQL final : public KVStore {
bool saveAll() override;

private:
std::vector<std::string> loadPrefix(const std::string &prefix = "") override;
std::optional<ValueWrapper> load(const std::string &key) override;
bool save(const std::string &key, const ValueWrapper &value) override;
bool prepareSave(const std::string &key, const ValueWrapper &value, DBInsert &update);
Expand Down
17 changes: 9 additions & 8 deletions src/kv/value_wrapper.cpp
Original file line number Diff line number Diff line change
@@ -1,30 +1,31 @@
#include "kv/value_wrapper.hpp"
#include "utils/tools.hpp"

ValueWrapper::ValueWrapper(uint64_t timestamp) :
timestamp_(timestamp) { }
timestamp_(timestamp == 0 ? getTimeMsNow() : timestamp) { }

ValueWrapper::ValueWrapper(const ValueVariant &value, uint64_t timestamp) :
data_(value), timestamp_(timestamp) { }
data_(value), timestamp_(timestamp == 0 ? getTimeMsNow() : timestamp) { }

ValueWrapper::ValueWrapper(const std::string &value, uint64_t timestamp) :
data_(value), timestamp_(timestamp) { }
data_(value), timestamp_(timestamp == 0 ? getTimeMsNow() : timestamp) { }

ValueWrapper::ValueWrapper(bool value, uint64_t timestamp) :
data_(value), timestamp_(timestamp) { }
data_(value), timestamp_(timestamp == 0 ? getTimeMsNow() : timestamp) { }

ValueWrapper::ValueWrapper(int value, uint64_t timestamp) :
data_(value), timestamp_(timestamp) { }
data_(value), timestamp_(timestamp == 0 ? getTimeMsNow() : timestamp) { }

ValueWrapper::ValueWrapper(double value, uint64_t timestamp) :
data_(value), timestamp_(timestamp) { }
data_(value), timestamp_(timestamp == 0 ? getTimeMsNow() : timestamp) { }

ValueWrapper::ValueWrapper(const phmap::flat_hash_map<std::string, ValueWrapper> &value, uint64_t timestamp) :
data_(createMapFromRange(value.begin(), value.end(), timestamp)),
timestamp_(timestamp) { }
timestamp_(timestamp == 0 ? getTimeMsNow() : timestamp) { }

ValueWrapper::ValueWrapper(const std::initializer_list<std::pair<const std::string, ValueWrapper>> &init_list, uint64_t timestamp) :
data_(createMapFromRange(init_list.begin(), init_list.end(), timestamp)),
timestamp_(timestamp) { }
timestamp_(timestamp == 0 ? getTimeMsNow() : timestamp) { }

std::optional<ValueWrapper> ValueWrapper::get(const std::string &key) const {
auto pval = std::get_if<MapType>(&data_);
Expand Down
40 changes: 39 additions & 1 deletion src/lua/functions/core/libs/kv_functions.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -72,14 +72,52 @@ int KVFunctions::luaKVGet(lua_State* L) {
valueWrapper = g_kv().get(key, forceLoad);
}

if (valueWrapper) {
if (valueWrapper.has_value()) {
pushValueWrapper(L, *valueWrapper);
} else {
lua_pushnil(L);
}
return 1;
}

int KVFunctions::luaKVRemove(lua_State* L) {
// KV.remove(key) | scopedKV:remove(key)
auto key = getString(L, -1);
if (isUserdata(L, 1)) {
auto scopedKV = getUserdataShared<KV>(L, 1);
scopedKV->remove(key);
} else {
g_kv().remove(key);
}
lua_pushnil(L);
return 1;
}

int KVFunctions::luaKVKeys(lua_State* L) {
// KV.keys([prefix = ""]) | scopedKV:keys([prefix = ""])
std::unordered_set<std::string> keys;
std::string prefix = "";

if (isString(L, -1)) {
prefix = getString(L, -1);
}

if (isUserdata(L, 1)) {
auto scopedKV = getUserdataShared<KV>(L, 1);
keys = scopedKV->keys();
} else {
keys = g_kv().keys(prefix);
}

int index = 0;
lua_createtable(L, static_cast<int>(keys.size()), 0);
for (const auto &key : keys) {
pushString(L, key);
lua_rawseti(L, -2, ++index);
}
return 1;
}

std::optional<ValueWrapper> KVFunctions::getValueWrapper(lua_State* L) {
if (isBoolean(L, -1)) {
return ValueWrapper(getBoolean(L, -1));
Expand Down
6 changes: 6 additions & 0 deletions src/lua/functions/core/libs/kv_functions.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -18,17 +18,23 @@ class KVFunctions final : LuaScriptInterface {
registerMethod(L, "kv", "scoped", KVFunctions::luaKVScoped);
registerMethod(L, "kv", "set", KVFunctions::luaKVSet);
registerMethod(L, "kv", "get", KVFunctions::luaKVGet);
registerMethod(L, "kv", "keys", KVFunctions::luaKVKeys);
registerMethod(L, "kv", "remove", KVFunctions::luaKVRemove);

registerClass(L, "KV", "");
registerMethod(L, "KV", "scoped", KVFunctions::luaKVScoped);
registerMethod(L, "KV", "set", KVFunctions::luaKVSet);
registerMethod(L, "KV", "get", KVFunctions::luaKVGet);
registerMethod(L, "KV", "keys", KVFunctions::luaKVKeys);
registerMethod(L, "KV", "remove", KVFunctions::luaKVRemove);
}

private:
static int luaKVScoped(lua_State* L);
static int luaKVSet(lua_State* L);
static int luaKVGet(lua_State* L);
static int luaKVKeys(lua_State* L);
static int luaKVRemove(lua_State* L);

static std::optional<ValueWrapper> getValueWrapper(lua_State* L);
static void pushStringValue(lua_State* L, const StringType &value);
Expand Down
3 changes: 3 additions & 0 deletions tests/fixture/kv/in_memory_kv.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,9 @@ class KVMemory final : public KVStore {
}

protected:
std::vector<std::string> loadPrefix(const std::string &prefix = "") override {
return {};
}
std::optional<ValueWrapper> load(const std::string &key) override {
return std::nullopt;
}
Expand Down
Loading