Allow converting of learn cache from sqlite to redis

2 files changed, 66 insertions, 9 deletions

diff --git a/src/rspamadm/stat_convert.c b/src/rspamadm/stat_convert.c
index 36c2d69cc..15f21f567 100644
--- a/src/rspamadm/stat_convert.c
+++ b/src/rspamadm/stat_convert.c
@@ -30,6 +30,7 @@
 static gchar *source_db = NULL;
 static gchar *redis_host = NULL;
 static gchar *symbol = NULL;
+static gchar *cache_db = NULL;
 
 static void rspamadm_statconvert (gint argc, gchar **argv);
 static const char *rspamadm_statconvert_help (gboolean full_help);
@@ -44,6 +45,8 @@ struct rspamadm_command statconvert_command = {
 static GOptionEntry entries[] = {
 		{"database", 'd', 0, G_OPTION_ARG_FILENAME, &source_db,
 				"Input sqlite",      NULL},
+		{"cache", 'c', 0, G_OPTION_ARG_FILENAME, &cache_db,
+				"Input learn cache",      NULL},
 		{"host", 'h', 0, G_OPTION_ARG_STRING, &redis_host,
 				"Output redis ip (in format ip:port)", NULL},
 		{"symbol", 's', 0, G_OPTION_ARG_STRING, &symbol,
@@ -63,7 +66,8 @@ rspamadm_statconvert_help (gboolean full_help)
 				"Where options are:\n\n"
 				"-d: input sqlite\n"
 				"-h: output redis ip (in format ip:port)\n"
-				"-s: symbol in redis (e.g. BAYES_SPAM)\n";
+				"-s: symbol in redis (e.g. BAYES_SPAM)\n"
+				"-c: also convert data from the learn cache\n";
 	}
 	else {
 		help_str = "Convert statistics from sqlite3 to redis";
@@ -119,6 +123,11 @@ rspamadm_statconvert (gint argc, gchar **argv)
 	ucl_object_insert_key (obj, ucl_object_fromstring (symbol),
 			"symbol", 0, false);
 
+	if (cache_db != NULL) {
+		ucl_object_insert_key (obj, ucl_object_fromstring (cache_db),
+				"cache_db", 0, false);
+	}
+
 	rspamadm_execute_lua_ucl_subr (L,
 			argc,
 			argv,
diff --git a/src/rspamadm/stat_convert.lua b/src/rspamadm/stat_convert.lua
index d33a22866..7a64a91e9 100644
--- a/src/rspamadm/stat_convert.lua
+++ b/src/rspamadm/stat_convert.lua
@@ -1,5 +1,6 @@
 local sqlite3 = require "rspamd_sqlite3"
 local redis = require "rspamd_redis"
+local util = require "rspamd_util"
 local _ = require "fun"
 
 local function send_redis(server, symbol, tokens)
@@ -18,7 +19,7 @@ local function send_redis(server, symbol, tokens)
       ret = false
     end
   end, tokens)
-  
+
   if ret then
     ret = conn:exec()
   end
@@ -26,6 +27,46 @@ local function send_redis(server, symbol, tokens)
   return ret
 end
 
+local function convert_learned(cache, target)
+  local converted = 0
+  local db = sqlite3.open(cache)
+  local ret = true
+
+  if not db then
+    print('Cannot open cache database: ' .. cache)
+    return false
+  end
+
+  db:sql('BEGIN;')
+  for row in db:rows('SELECT * FROM learns;') do
+    local is_spam
+    local digest = tostring(util.encode_base32(row.digest))
+
+    if row.flag == '0' then
+      is_spam = '-1'
+    else
+      is_spam = '1'
+    end
+
+    if not redis.make_request_sync({
+      host = target,
+      cmd = 'HSET',
+      args = {'learned_ids', digest, is_spam}
+    }) then
+      print('Cannot add hash: ' .. digest)
+      ret = false
+    else
+      converted = converted + 1
+    end
+  end
+  db:sql('COMMIT;')
+
+  print(string.format('Converted %d cached items from sqlite3 learned cache to redis',
+    converted))
+
+  return ret
+end
+
 return function (args, res)
   local db = sqlite3.open(res['source_db'])
   local tokens = {}
@@ -36,6 +77,13 @@ return function (args, res)
   local users_map = {}
   local learns = {}
 
+  if res['cache_db'] then
+    if not convert_learned(res['cache_db'], res['redis_host']) then
+      print('Cannot convert learned cache to redis')
+      return
+    end
+  end
+
   if not db then
     print('Cannot open source db: ' .. res['source_db'])
     return
@@ -58,9 +106,9 @@ return function (args, res)
     if row.user ~= 0 and users_map[row.user] then
       user = users_map[row.user]
     end
-    
+
     table.insert(tokens, {row.token, row.value, user})
-    
+
     num = num + 1
     total = total + 1
     if num > lim then
@@ -68,14 +116,14 @@ return function (args, res)
         print('Cannot send tokens to the redis server')
         return
       end
-      
+
       num = 0
       tokens = {}
     end
   end
-  if #tokens > 0 and 
+  if #tokens > 0 and
     not send_redis(res['redis_host'], res['symbol'], tokens, users_map) then
-    
+
     print('Cannot send tokens to the redis server')
     return
   end
@@ -91,7 +139,7 @@ return function (args, res)
     end
   end, learns)
   db:sql('COMMIT;')
-  
+
   print(string.format('Migrated %d tokens for %d users for symbol %s',
     total, nusers, res['symbol']))
-end
\ No newline at end of file
+end