From 26f64fdb489893a898020d77c628ea3f881c81ab Mon Sep 17 00:00:00 2001 From: Vsevolod Stakhov Date: Mon, 25 Jan 2016 20:13:48 +0000 Subject: [PATCH] Add stat_convert command New command is intended to convert sqlite stats to redis stats --- src/rspamadm/CMakeLists.txt | 1 + src/rspamadm/commands.c | 2 + src/rspamadm/stat_convert.c | 130 ++++++++++++++++++++++++++++++++++ src/rspamadm/stat_convert.lua | 70 ++++++++++++++++++ 4 files changed, 203 insertions(+) create mode 100644 src/rspamadm/stat_convert.c create mode 100644 src/rspamadm/stat_convert.lua diff --git a/src/rspamadm/CMakeLists.txt b/src/rspamadm/CMakeLists.txt index 1fc801b46..7d5170cf5 100644 --- a/src/rspamadm/CMakeLists.txt +++ b/src/rspamadm/CMakeLists.txt @@ -7,6 +7,7 @@ SET(RSPAMADMSRC rspamadm.c configdump.c control.c confighelp.c + stat_convert.c ${CMAKE_BINARY_DIR}/src/workers.c ${CMAKE_BINARY_DIR}/src/modules.c ${CMAKE_SOURCE_DIR}/src/controller.c diff --git a/src/rspamadm/commands.c b/src/rspamadm/commands.c index 589e15440..4514376e1 100644 --- a/src/rspamadm/commands.c +++ b/src/rspamadm/commands.c @@ -30,6 +30,7 @@ extern struct rspamadm_command fuzzy_merge_command; extern struct rspamadm_command configdump_command; extern struct rspamadm_command control_command; extern struct rspamadm_command confighelp_command; +extern struct rspamadm_command statconvert_command; const struct rspamadm_command *commands[] = { &help_command, @@ -40,6 +41,7 @@ const struct rspamadm_command *commands[] = { &configdump_command, &control_command, &confighelp_command, + &statconvert_command, NULL }; diff --git a/src/rspamadm/stat_convert.c b/src/rspamadm/stat_convert.c new file mode 100644 index 000000000..36c2d69cc --- /dev/null +++ b/src/rspamadm/stat_convert.c @@ -0,0 +1,130 @@ +/* + * Copyright (c) 2016, Vsevolod Stakhov + * All rights reserved. + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright + * notice, this list of conditions and the following disclaimer in the + * documentation and/or other materials provided with the distribution. + * + * THIS SOFTWARE IS PROVIDED BY AUTHOR ''AS IS'' AND ANY + * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED + * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL AUTHOR BE LIABLE FOR ANY + * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES + * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; + * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND + * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS + * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + */ + +#include "config.h" +#include "rspamadm.h" +#include "lua/lua_common.h" +#include "stat_convert.lua.h" + +static gchar *source_db = NULL; +static gchar *redis_host = NULL; +static gchar *symbol = NULL; + +static void rspamadm_statconvert (gint argc, gchar **argv); +static const char *rspamadm_statconvert_help (gboolean full_help); + +struct rspamadm_command statconvert_command = { + .name = "statconvert", + .flags = 0, + .help = rspamadm_statconvert_help, + .run = rspamadm_statconvert +}; + +static GOptionEntry entries[] = { + {"database", 'd', 0, G_OPTION_ARG_FILENAME, &source_db, + "Input sqlite", 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, + "Symbol in redis (e.g. BAYES_SPAM)", NULL}, + {NULL, 0, 0, G_OPTION_ARG_NONE, NULL, NULL, NULL} +}; + + +static const char * +rspamadm_statconvert_help (gboolean full_help) +{ + const char *help_str; + + if (full_help) { + help_str = "Convert statistics from sqlite3 to redis\n\n" + "Usage: rspamadm statconvert -d -h -s \n" + "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"; + } + else { + help_str = "Convert statistics from sqlite3 to redis"; + } + + return help_str; +} + +static void +rspamadm_statconvert (gint argc, gchar **argv) +{ + GOptionContext *context; + GError *error = NULL; + lua_State *L; + ucl_object_t *obj; + + context = g_option_context_new ( + "statconvert - converts statistics from sqlite3 to redis"); + g_option_context_set_summary (context, + "Summary:\n Rspamd administration utility version " + RVERSION + "\n Release id: " + RID); + g_option_context_add_main_entries (context, entries, NULL); + g_option_context_set_ignore_unknown_options (context, TRUE); + + if (!g_option_context_parse (context, &argc, &argv, &error)) { + rspamd_fprintf (stderr, "option parsing failed: %s\n", error->message); + g_error_free (error); + exit (1); + } + + if (!source_db) { + rspamd_fprintf (stderr, "source db is missing\n"); + exit (1); + } + if (!redis_host) { + rspamd_fprintf (stderr, "redis host is missing\n"); + exit (1); + } + if (!symbol) { + rspamd_fprintf (stderr, "symbol is missing\n"); + exit (1); + } + + L = rspamd_lua_init (); + + obj = ucl_object_typed_new (UCL_OBJECT); + ucl_object_insert_key (obj, ucl_object_fromstring (source_db), + "source_db", 0, false); + ucl_object_insert_key (obj, ucl_object_fromstring (redis_host), + "redis_host", 0, false); + ucl_object_insert_key (obj, ucl_object_fromstring (symbol), + "symbol", 0, false); + + rspamadm_execute_lua_ucl_subr (L, + argc, + argv, + obj, + rspamadm_script_stat_convert); + + lua_close (L); + ucl_object_unref (obj); +} diff --git a/src/rspamadm/stat_convert.lua b/src/rspamadm/stat_convert.lua new file mode 100644 index 000000000..05e93fc49 --- /dev/null +++ b/src/rspamadm/stat_convert.lua @@ -0,0 +1,70 @@ +local sqlite3 = require "rspamd_sqlite3" +local redis = require "rspamd_redis" +local _ = require "fun" + +local function send_redis(server, symbol, tokens) + local ret = true + local args = {} + + _.each(function(t) + if not args[t[3]] then + args[t[3]] = {symbol .. t[3]} + end + table.insert(args[t[3]], t[1]) + table.insert(args[t[3]], t[2]) + end, tokens) + + _.each(function(k, argv) + if not redis.make_request_sync({ + host = server, + cmd = 'HMSET', + args = argv + }) then + ret = false + end + end, args) + + return ret +end + +return function (args, res) + local db = sqlite3.open(res['source_db']) + local tokens = {} + local num = 0 + local lim = 100 -- Update each 100 tokens + local users_map = {} + local learns = {} + + if not db then + print('Cannot open source db: ' .. res['source_db']) + return + end + + db:sql('BEGIN;') + -- Fill users mapping + for row in db:rows('SELECT * FROM users;') do + users_map[row.id] = row.name + learns[row.id] = row.learned + end + -- Fill tokens, sending data to redis each `lim` records + for row in db:rows('SELECT token,value,user FROM tokens;') do + local user = '' + 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 + if num > lim then + if not send_redis(res['redis_host'], res['symbol'], tokens, users_map) then + print('Cannot send tokens to the redis server') + return + end + + num = 0 + tokens = {} + end + end + db:sql('COMMIT;') +end \ No newline at end of file -- 2.39.5