[Feature] Allow to reset redis tokens instead of appendig values

2 files changed, 17 insertions, 6 deletions

diff --git a/src/rspamadm/stat_convert.c b/src/rspamadm/stat_convert.c
index ae4b00ecf..48acbc9bd 100644
--- a/src/rspamadm/stat_convert.c
+++ b/src/rspamadm/stat_convert.c
@@ -24,6 +24,7 @@ static gchar *symbol = NULL;
 static gchar *cache_db = NULL;
 static gchar *redis_db = NULL;
 static gchar *redis_password = NULL;
+static gboolean reset_previous = FALSE;
 
 static void rspamadm_statconvert (gint argc, gchar **argv);
 static const char *rspamadm_statconvert_help (gboolean full_help);
@@ -48,6 +49,8 @@ static GOptionEntry entries[] = {
 				"Database in redis (should be numeric)", NULL},
 		{"password", 'p', 0, G_OPTION_ARG_STRING, &redis_password,
 				"Password to connect to redis", NULL},
+		{"reset", 'r', 0, G_OPTION_ARG_NONE, &reset_previous,
+				"Reset previous data instead of appending values", NULL},
 		{NULL,     0,   0, G_OPTION_ARG_NONE, NULL, NULL, NULL}
 };
 
@@ -66,7 +69,8 @@ rspamadm_statconvert_help (gboolean full_help)
 				"-s: symbol in redis (e.g. BAYES_SPAM)\n"
 				"-c: also convert data from the learn cache\n"
 				"-D: output redis database\n"
-				"-p: redis password\n";
+				"-p: redis password\n"
+				"-r: reset previous data instead of increasing values\n";
 	}
 	else {
 		help_str = "Convert statistics from sqlite3 to redis";
@@ -121,6 +125,8 @@ rspamadm_statconvert (gint argc, gchar **argv)
 			"redis_host", 0, false);
 	ucl_object_insert_key (obj, ucl_object_fromstring (symbol),
 			"symbol", 0, false);
+	ucl_object_insert_key (obj, ucl_object_frombool (reset_previous),
+			"reset_previous", 0, false);
 
 	if (cache_db != NULL) {
 		ucl_object_insert_key (obj, ucl_object_fromstring (cache_db),
diff --git a/src/rspamadm/stat_convert.lua b/src/rspamadm/stat_convert.lua
index fa03a6215..c0cb4e0eb 100644
--- a/src/rspamadm/stat_convert.lua
+++ b/src/rspamadm/stat_convert.lua
@@ -2,7 +2,7 @@ local sqlite3 = require "rspamd_sqlite3"
 local redis = require "rspamd_redis"
 local util = require "rspamd_util"
 
-local function send_redis(server, symbol, tokens, password, db)
+local function send_redis(server, symbol, tokens, password, db, cmd)
   local ret = true
   local conn,err = redis.connect_sync({
     host = server,
@@ -21,7 +21,7 @@ local function send_redis(server, symbol, tokens, password, db)
   end
 
   for _,t in ipairs(tokens) do
-    if not conn:add_cmd('HINCRBY', {symbol .. t[3], t[1], t[2]}) then
+    if not conn:add_cmd(cmd, {symbol .. t[3], t[1], t[2]}) then
       ret = false
     end
   end
@@ -106,6 +106,11 @@ return function (args, res)
   local redis_password = res['redis_password']
   local redis_db = res['redis_db']
   local ret = false
+  local cmd = 'HINCRBY'
+
+  if res['reset_previous'] then
+    cmd = 'HSET'
+  end
 
   if res['cache_db'] then
     if not convert_learned(res['cache_db'], res['redis_host']) then
@@ -153,7 +158,7 @@ return function (args, res)
     total = total + 1
     if num > lim then
       if not send_redis(res['redis_host'], res['symbol'],
-        tokens, redis_password, redis_db) then
+        tokens, redis_password, redis_db, cmd) then
 
         print('Cannot send tokens to the redis server')
         return
@@ -165,7 +170,7 @@ return function (args, res)
   end
   if #tokens > 0 and
     not send_redis(res['redis_host'], res['symbol'], tokens,
-      redis_password, redis_db) then
+      redis_password, redis_db, cmd) then
 
     print('Cannot send tokens to the redis server')
     return
@@ -189,7 +194,7 @@ return function (args, res)
 
   for id,learned in pairs(learns) do
     local user = users_map[id]
-    if not conn:add_cmd('HINCRBY', {res['symbol'] .. user, 'learns', learned}) then
+    if not conn:add_cmd(cmd, {res['symbol'] .. user, 'learns', learned}) then
       print('Cannot update learns for user: ' .. user)
     end
   end