]> source.dussan.org Git - rspamd.git/commitdiff
Add languages support to sqlite3 backend.
authorVsevolod Stakhov <vsevolod@highsecure.ru>
Thu, 23 Jul 2015 18:14:04 +0000 (19:14 +0100)
committerVsevolod Stakhov <vsevolod@highsecure.ru>
Thu, 23 Jul 2015 18:14:04 +0000 (19:14 +0100)
src/libstat/backends/sqlite3_backend.c

index 0852bcf5e0ce4cb20c25a645c48e5a9703540049..af84faef0411802f93bda4ea54241bef53ffb1b8 100644 (file)
@@ -27,6 +27,7 @@
 #include "sqlite3.h"
 #include "libutil/sqlite_utils.h"
 #include "libstat/stat_internal.h"
+#include "libmime/message.h"
 
 #define SQLITE3_BACKEND_TYPE "sqlite3"
 #define SQLITE3_SCHEMA_VERSION "1"
@@ -154,8 +155,8 @@ static struct rspamd_sqlite3_prstmt prepared_stmts[RSPAMD_STAT_BACKEND_MAX] =
        },
        {
                .idx = RSPAMD_STAT_BACKEND_INC_LEARNS,
-               .sql = "UPDATE languages SET learns=learns + 1 WHERE name=?1;"
-                               "UPDATE users SET learns=learns + 1 WHERE name=?2;",
+               .sql = "UPDATE languages SET learns=learns + 1 WHERE id=?1;"
+                               "UPDATE users SET learns=learns + 1 WHERE id=?2;",
                .stmt = NULL,
                .args = "TT",
                .result = SQLITE_DONE,
@@ -163,8 +164,8 @@ static struct rspamd_sqlite3_prstmt prepared_stmts[RSPAMD_STAT_BACKEND_MAX] =
        },
        {
                .idx = RSPAMD_STAT_BACKEND_DEC_LEARNS,
-               .sql = "UPDATE languages SET learns=learns - 1 WHERE name=?1;"
-                               "UPDATE users SET learns=learns - 1 WHERE name=?2;",
+               .sql = "UPDATE languages SET learns=learns - 1 WHERE id=?1;"
+                               "UPDATE users SET learns=learns - 1 WHERE id=?2;",
                .stmt = NULL,
                .args = "TT",
                .result = SQLITE_DONE,
@@ -268,6 +269,47 @@ rspamd_sqlite3_get_user (struct rspamd_stat_sqlite3_db *db,
        return id;
 }
 
+static gint64
+rspamd_sqlite3_get_language (struct rspamd_stat_sqlite3_db *db,
+               struct rspamd_task *task, gboolean learn)
+{
+       gint64 id = 0; /* Default language is 0 */
+       gint rc;
+       guint i;
+       const gchar *language = NULL;
+       struct mime_text_part *tp;
+
+       for (i = 0; i < task->text_parts->len; i ++) {
+               tp = g_ptr_array_index (task->text_parts, i);
+
+               if (tp->lang_code != NULL && tp->lang_code[0] != '\0' &&
+                               strcmp (tp->lang_code, "en") != 0) {
+                       language = tp->language;
+                       break;
+               }
+       }
+
+       /* XXX: We ignore multiple languages but default + extra */
+       if (language != NULL) {
+               rc = rspamd_sqlite3_run_prstmt (db->sqlite, db->prstmt,
+                               RSPAMD_STAT_BACKEND_GET_LANGUAGE, language, &id);
+
+               if (rc != SQLITE_OK && learn) {
+                       /* We need to insert a new language */
+                       if (!db->in_transaction) {
+                               rspamd_sqlite3_run_prstmt (db->sqlite, db->prstmt,
+                                               RSPAMD_STAT_BACKEND_TRANSACTION_START_IM);
+                               db->in_transaction = TRUE;
+                       }
+
+                       rc =  rspamd_sqlite3_run_prstmt (db->sqlite, db->prstmt,
+                                       RSPAMD_STAT_BACKEND_INSERT_LANGUAGE, language, &id);
+               }
+       }
+
+       return id;
+}
+
 static struct rspamd_stat_sqlite3_db *
 rspamd_sqlite3_opendb (const gchar *path, const ucl_object_t *opts,
                gboolean create, GError **err)
@@ -402,11 +444,12 @@ rspamd_sqlite3_runtime (struct rspamd_task *task,
        bk = g_hash_table_lookup (ctx->files, stcf);
 
        if (bk) {
-               rt = rspamd_mempool_alloc0 (task->task_pool, sizeof (*rt));
+               rt = rspamd_mempool_alloc (task->task_pool, sizeof (*rt));
                rt->ctx = ctx;
                rt->db = bk;
                rt->task = task;
                rt->user_id = -1;
+               rt->lang_id = -1;
        }
 
        return rt;
@@ -445,6 +488,10 @@ rspamd_sqlite3_process_token (struct rspamd_task *task, struct token_node_s *tok
                rt->user_id = rspamd_sqlite3_get_user (bk, task, FALSE);
        }
 
+       if (rt->lang_id == -1) {
+               rt->lang_id = rspamd_sqlite3_get_language (bk, task, FALSE);
+       }
+
        memcpy (&idx, tok->data, sizeof (idx));
 
        /* TODO: language and user support */
@@ -518,6 +565,10 @@ rspamd_sqlite3_learn_token (struct rspamd_task *task, struct token_node_s *tok,
                rt->user_id = rspamd_sqlite3_get_user (bk, task, TRUE);
        }
 
+       if (rt->lang_id == -1) {
+               rt->lang_id = rspamd_sqlite3_get_language (bk, task, FALSE);
+       }
+
        iv = res->value;
        memcpy (&idx, tok->data, sizeof (idx));
 
@@ -577,7 +628,7 @@ rspamd_sqlite3_inc_learns (struct rspamd_task *task, gpointer runtime,
        bk = rt->db;
        rspamd_sqlite3_run_prstmt (bk->sqlite, bk->prstmt,
                        RSPAMD_STAT_BACKEND_INC_LEARNS,
-                       SQLITE3_DEFAULT, SQLITE3_DEFAULT);
+                       rt->user_id, rt->lang_id);
 
        if (bk->in_transaction) {
                rspamd_sqlite3_run_prstmt (bk->sqlite, bk->prstmt,
@@ -603,7 +654,7 @@ rspamd_sqlite3_dec_learns (struct rspamd_task *task, gpointer runtime,
        bk = rt->db;
        rspamd_sqlite3_run_prstmt (bk->sqlite, bk->prstmt,
                        RSPAMD_STAT_BACKEND_DEC_LEARNS,
-                       SQLITE3_DEFAULT, SQLITE3_DEFAULT);
+                       rt->user_id, rt->lang_id);
 
        if (bk->in_transaction) {
                rspamd_sqlite3_run_prstmt (bk->sqlite, bk->prstmt,