From 4ce6722cedb43293ee58e96b84c40dd9ed4b7a9c Mon Sep 17 00:00:00 2001 From: Vsevolod Stakhov Date: Wed, 16 Mar 2016 16:42:43 +0000 Subject: [PATCH] [Feature] Allow to set db and password for redis in stat_convert --- src/rspamadm/stat_convert.c | 20 ++++++- src/rspamadm/stat_convert.lua | 104 +++++++++++++++++++++++++++------- 2 files changed, 102 insertions(+), 22 deletions(-) diff --git a/src/rspamadm/stat_convert.c b/src/rspamadm/stat_convert.c index 0a6617074..ae4b00ecf 100644 --- 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 7f24890d0..fa03a6215 100644 --- 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 -- 2.39.5