[Feature] Allow to set db and password for redis in stat_convert

diff --git a/src/rspamadm/stat_convert.c b/src/rspamadm/stat_convert.c
index 0a6617074d65b0c1a71ccb2bd5c10d7b0af58d2f..ae4b00ecf9ee08c45be90b99e41bc088ba4f86aa 100644 (file)
--- a/src/rspamadm/stat_convert.c
+++ b/src/rspamadm/stat_convert.c
@@ -22,6 +22,8 @@ static gchar *source_db = NULL;
 static gchar *redis_host = NULL;
 static gchar *symbol = NULL;
 static gchar *cache_db = NULL;
+static gchar *redis_db = NULL;
+static gchar *redis_password = NULL;
 
 static void rspamadm_statconvert (gint argc, gchar **argv);
 static const char *rspamadm_statconvert_help (gboolean full_help);
@@ -42,6 +44,10 @@ static GOptionEntry entries[] = {
                                "Output redis ip (in format ip:port)", NULL},
                {"symbol", 's', 0, G_OPTION_ARG_STRING, &symbol,
                                "Symbol in redis (e.g. BAYES_SPAM)", NULL},
+               {"dbname", 'D', 0, G_OPTION_ARG_STRING, &redis_db,
+                               "Database in redis (should be numeric)", NULL},
+               {"password", 'p', 0, G_OPTION_ARG_STRING, &redis_password,
+                               "Password to connect to redis", NULL},
                {NULL,     0,   0, G_OPTION_ARG_NONE, NULL, NULL, NULL}
 };
 
@@ -58,7 +64,9 @@ rspamadm_statconvert_help (gboolean full_help)
                                "-d: input sqlite\n"
                                "-h: output redis ip (in format ip:port)\n"
                                "-s: symbol in redis (e.g. BAYES_SPAM)\n"
-                               "-c: also convert data from the learn cache\n";
+                               "-c: also convert data from the learn cache\n"
+                               "-D: output redis database\n"
+                               "-p: redis password\n";
        }
        else {
                help_str = "Convert statistics from sqlite3 to redis";
@@ -119,6 +127,16 @@ rspamadm_statconvert (gint argc, gchar **argv)
                                "cache_db", 0, false);
        }
 
+       if (redis_password) {
+               ucl_object_insert_key (obj, ucl_object_fromstring (redis_password),
+                               "redis_password", 0, false);
+       }
+
+       if (redis_db) {
+               ucl_object_insert_key (obj, ucl_object_fromstring (redis_db),
+                               "redis_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 7f24890d0a09c48bab2b94d0bb6ea47e0cf4219d..fa03a621599b534a076ddcd86a6c598b50b8df96 100644 (file)
--- a/src/rspamadm/stat_convert.lua
+++ b/src/rspamadm/stat_convert.lua
@@ -2,17 +2,24 @@ local sqlite3 = require "rspamd_sqlite3"
 local redis = require "rspamd_redis"
 local util = require "rspamd_util"
 
-local function send_redis(server, symbol, tokens)
+local function send_redis(server, symbol, tokens, password, db)
   local ret = true
-  local conn = redis.connect_sync({
+  local conn,err = redis.connect_sync({
     host = server,
   })
 
   if not conn then
-    print('Cannot connect to ' .. server)
+    print('Cannot connect to ' .. server .. ' error: ' .. err)
     return false
   end
 
+  if password then
+    conn:add_cmd('AUTH', {password})
+  end
+  if db then
+    conn:add_cmd('SELECT', {db})
+  end
+
   for _,t in ipairs(tokens) do
     if not conn:add_cmd('HINCRBY', {symbol .. t[3], t[1], t[2]}) then
       ret = false
@@ -26,7 +33,7 @@ local function send_redis(server, symbol, tokens)
   return ret
 end
 
-local function convert_learned(cache, target)
+local function convert_learned(cache, server, password, db)
   local converted = 0
   local db = sqlite3.open(cache)
   local ret = true
@@ -37,6 +44,23 @@ local function convert_learned(cache, target)
   end
 
   db:sql('BEGIN;')
+
+  local conn,err = redis.connect_sync({
+    host = server,
+  })
+
+  if not conn then
+    print('Cannot connect to ' .. server .. ' error: ' .. err)
+    return false
+  end
+
+  if password then
+    conn:add_cmd('AUTH', {password})
+  end
+  if db then
+    conn:add_cmd('SELECT', {db})
+  end
+
   for row in db:rows('SELECT * FROM learns;') do
     local is_spam
     local digest = tostring(util.encode_base32(row.digest))
@@ -47,11 +71,7 @@ local function convert_learned(cache, target)
       is_spam = '1'
     end
 
-    if not redis.make_request_sync({
-      host = target,
-      cmd = 'HSET',
-      args = {'learned_ids', digest, is_spam}
-    }) then
+    if not conn:add_cmd('HSET', {'learned_ids', digest, is_spam}) then
       print('Cannot add hash: ' .. digest)
       ret = false
     else
@@ -60,8 +80,16 @@ local function convert_learned(cache, target)
   end
   db:sql('COMMIT;')
 
-  print(string.format('Converted %d cached items from sqlite3 learned cache to redis',
-    converted))
+  if ret then
+    ret = conn:exec()
+  end
+
+  if ret then
+    print(string.format('Converted %d cached items from sqlite3 learned cache to redis',
+      converted))
+  else
+    print('Error occurred during sending data to redis')
+  end
 
   return ret
 end
@@ -75,6 +103,9 @@ return function (args, res)
   local lim = 1000 -- Update each 1000 tokens
   local users_map = {}
   local learns = {}
+  local redis_password = res['redis_password']
+  local redis_db = res['redis_db']
+  local ret = false
 
   if res['cache_db'] then
     if not convert_learned(res['cache_db'], res['redis_host']) then
@@ -96,9 +127,19 @@ return function (args, res)
     else
       users_map[row.id] = row.name
     end
-    learns[row.id] = row.learned
+    learns[row.id] = row.learns
     nusers = nusers + 1
   end
+
+  -- Workaround for old databases
+  for row in db:rows('SELECT * FROM languages WHERE id=0;') do
+    if learns[row.id] then
+      learns[row.id] = learns[row.id] + row.learns
+    else
+      learns[row.id] = row.learns
+    end
+  end
+
   -- Fill tokens, sending data to redis each `lim` records
   for row in db:rows('SELECT token,value,user FROM tokens;') do
     local user = ''
@@ -111,7 +152,9 @@ return function (args, res)
     num = num + 1
     total = total + 1
     if num > lim then
-      if not send_redis(res['redis_host'], res['symbol'], tokens, users_map) then
+      if not send_redis(res['redis_host'], res['symbol'],
+        tokens, redis_password, redis_db) then
+
         print('Cannot send tokens to the redis server')
         return
       end
@@ -121,24 +164,43 @@ return function (args, res)
     end
   end
   if #tokens > 0 and
-    not send_redis(res['redis_host'], res['symbol'], tokens, users_map) then
+    not send_redis(res['redis_host'], res['symbol'], tokens,
+      redis_password, redis_db) then
 
     print('Cannot send tokens to the redis server')
     return
   end
   -- Now update all users
+  local conn,err = redis.connect_sync({
+    host = res['redis_host'],
+  })
+
+  if not conn then
+    print('Cannot connect to ' .. res['redis_host'] .. ' error: ' .. err)
+    return false
+  end
+
+  if redis_password then
+    conn:add_cmd('AUTH', {redis_password})
+  end
+  if redis_db then
+    conn:add_cmd('SELECT', {redis_db})
+  end
+
   for id,learned in pairs(learns) do
     local user = users_map[id]
-    if not redis.make_request_sync({
-        host = server,
-        cmd = 'HSET',
-        args = {symbol .. user, 'learns', learned}
-      }) then
+    if not conn:add_cmd('HINCRBY', {res['symbol'] .. user, 'learns', learned}) then
       print('Cannot update learns for user: ' .. user)
     end
   end
   db:sql('COMMIT;')
 
-  print(string.format('Migrated %d tokens for %d users for symbol %s',
-    total, nusers, res['symbol']))
+  ret = conn:exec()
+
+  if ret then
+    print(string.format('Migrated %d tokens for %d users for symbol %s',
+     total, nusers, res['symbol']))
+  else
+    print('Error occurred during sending data to redis')
+  end
 end