diff --git a/src/kv/kv.cpp b/src/kv/kv.cpp index 0e2f4b2f836..f3133dec9ca 100644 --- a/src/kv/kv.cpp +++ b/src/kv/kv.cpp @@ -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; @@ -53,7 +53,7 @@ void KVStore::setLocked(const std::string &key, const ValueWrapper &value) { } std::optional 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); @@ -72,11 +72,26 @@ std::optional KVStore::get(const std::string &key, bool forceLoad return value; } +std::unordered_set KVStore::keys(const std::string &prefix /*= ""*/) { + std::scoped_lock lock(mutex_); + std::unordered_set 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 KVStore::scoped(const std::string &scope) { - logger.debug("KVStore::scoped({})", scope); + logger.trace("KVStore::scoped({})", scope); return std::make_shared(logger, *this, scope); } diff --git a/src/kv/kv.hpp b/src/kv/kv.hpp index 40fe449834f..99bc1f695a6 100644 --- a/src/kv/kv.hpp +++ b/src/kv/kv.hpp @@ -32,6 +32,8 @@ class KV : public std::enable_shared_from_this { virtual std::shared_ptr scoped(const std::string &scope) = 0; + virtual std::unordered_set keys(const std::string &prefix = "") = 0; + void remove(const std::string &key); virtual void flush() { @@ -60,6 +62,7 @@ class KVStore : public KV { } std::shared_ptr scoped(const std::string &scope) override final; + std::unordered_set keys(const std::string &prefix = ""); protected: phmap::parallel_flat_hash_map::iterator>> getStore() { @@ -76,6 +79,7 @@ class KVStore : public KV { virtual std::optional load(const std::string &key) = 0; virtual bool save(const std::string &key, const ValueWrapper &value) = 0; + virtual std::vector loadPrefix(const std::string &prefix = "") = 0; private: void setLocked(const std::string &key, const ValueWrapper &value); @@ -118,10 +122,14 @@ class ScopedKV final : public KV { } std::shared_ptr scoped(const std::string &scope) override final { - logger.debug("ScopedKV::scoped({})", buildKey(scope)); + logger.trace("ScopedKV::scoped({})", buildKey(scope)); return std::make_shared(logger, rootKV_, buildKey(scope)); } + std::unordered_set 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); diff --git a/src/kv/kv_sql.cpp b/src/kv/kv_sql.cpp index a0f1623ed2c..4784efa5c9e 100644 --- a/src/kv/kv_sql.cpp +++ b/src/kv/kv_sql.cpp @@ -41,6 +41,24 @@ std::optional KVSQL::load(const std::string &key) { return std::nullopt; } +std::vector KVSQL::loadPrefix(const std::string &prefix /* = ""*/) { + std::vector 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); @@ -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; } diff --git a/src/kv/kv_sql.hpp b/src/kv/kv_sql.hpp index 8cb4dce89f7..fd181b1cef7 100644 --- a/src/kv/kv_sql.hpp +++ b/src/kv/kv_sql.hpp @@ -25,6 +25,7 @@ class KVSQL final : public KVStore { bool saveAll() override; private: + std::vector loadPrefix(const std::string &prefix = "") override; std::optional 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); diff --git a/src/kv/value_wrapper.cpp b/src/kv/value_wrapper.cpp index 023beba1595..af304eeacad 100644 --- a/src/kv/value_wrapper.cpp +++ b/src/kv/value_wrapper.cpp @@ -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 &value, uint64_t timestamp) : data_(createMapFromRange(value.begin(), value.end(), timestamp)), - timestamp_(timestamp) { } + timestamp_(timestamp == 0 ? getTimeMsNow() : timestamp) { } ValueWrapper::ValueWrapper(const std::initializer_list> &init_list, uint64_t timestamp) : data_(createMapFromRange(init_list.begin(), init_list.end(), timestamp)), - timestamp_(timestamp) { } + timestamp_(timestamp == 0 ? getTimeMsNow() : timestamp) { } std::optional ValueWrapper::get(const std::string &key) const { auto pval = std::get_if(&data_); diff --git a/src/lua/functions/core/libs/kv_functions.cpp b/src/lua/functions/core/libs/kv_functions.cpp index 769c20bff8a..dbef916364d 100644 --- a/src/lua/functions/core/libs/kv_functions.cpp +++ b/src/lua/functions/core/libs/kv_functions.cpp @@ -72,7 +72,7 @@ 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); @@ -80,6 +80,44 @@ int KVFunctions::luaKVGet(lua_State* 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(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 keys; + std::string prefix = ""; + + if (isString(L, -1)) { + prefix = getString(L, -1); + } + + if (isUserdata(L, 1)) { + auto scopedKV = getUserdataShared(L, 1); + keys = scopedKV->keys(); + } else { + keys = g_kv().keys(prefix); + } + + int index = 0; + lua_createtable(L, static_cast(keys.size()), 0); + for (const auto &key : keys) { + pushString(L, key); + lua_rawseti(L, -2, ++index); + } + return 1; +} + std::optional KVFunctions::getValueWrapper(lua_State* L) { if (isBoolean(L, -1)) { return ValueWrapper(getBoolean(L, -1)); diff --git a/src/lua/functions/core/libs/kv_functions.hpp b/src/lua/functions/core/libs/kv_functions.hpp index 243a0ed36e3..3abddfb0b53 100644 --- a/src/lua/functions/core/libs/kv_functions.hpp +++ b/src/lua/functions/core/libs/kv_functions.hpp @@ -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 getValueWrapper(lua_State* L); static void pushStringValue(lua_State* L, const StringType &value); diff --git a/tests/fixture/kv/in_memory_kv.hpp b/tests/fixture/kv/in_memory_kv.hpp index b738a98470a..84a84ba8dff 100644 --- a/tests/fixture/kv/in_memory_kv.hpp +++ b/tests/fixture/kv/in_memory_kv.hpp @@ -35,6 +35,9 @@ class KVMemory final : public KVStore { } protected: + std::vector loadPrefix(const std::string &prefix = "") override { + return {}; + } std::optional load(const std::string &key) override { return std::nullopt; }