feat: kv: remove key and list prefix (#1994)

Adds a much needed `:remove()` method to the KV lua interface as well as a `:keys` method to list every key from a prefix. Fixes #1992 Related #1979
opentibiabr · Dec 10, 2023 · 5ebad34 · 5ebad34
1 parent c778b8e
commit 5ebad34
Show file tree

Hide file tree

Showing 8 changed files with 104 additions and 14 deletions.
diff --git 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<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);
@@ -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);
 }
diff --git a/src/kv/kv.hpp b/src/kv/kv.hpp
@@ -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() {
@@ -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() {
@@ -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);
@@ -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);

diff --git a/src/kv/kv_sql.cpp b/src/kv/kv_sql.cpp
@@ -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);
@@ -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
@@ -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);

diff --git 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<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_);

diff --git a/src/lua/functions/core/libs/kv_functions.cpp b/src/lua/functions/core/libs/kv_functions.cpp
@@ -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));

diff --git 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<ValueWrapper> 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
@@ -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;
 	}