commit 432064e: [Feature] Clickhouse: Store groups data

[Vsevolod Stakhov]

Author: Vsevolod Stakhov
Date: 2019-09-01 09:43:46 +0100
URL: https://github.com/rspamd/rspamd/commit/432064e1d636ec454851a70aef137cfd97e1cfee

[Feature] Clickhouse: Store groups data

---
 src/plugins/lua/clickhouse.lua                     | 33 +++++++++++++++++++++-
 .../cases/210_clickhouse/001_migration.robot       |  4 ++-
 2 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/src/plugins/lua/clickhouse.lua b/src/plugins/lua/clickhouse.lua
index 8762de239..5c8eb5d4d 100644
--- a/src/plugins/lua/clickhouse.lua
+++ b/src/plugins/lua/clickhouse.lua
@@ -31,7 +31,7 @@ end
 local data_rows = {}
 local custom_rows = {}
 local nrows = 0
-local schema_version = 7 -- Current schema version
+local schema_version = 8 -- Current schema version
 
 local settings = {
   limit = 1000,
@@ -131,6 +131,8 @@ CREATE TABLE rspamd
     `Symbols.Names` Array(LowCardinality(String)) COMMENT 'Symbol name',
     `Symbols.Scores` Array(Float32) COMMENT 'Symbol score',
     `Symbols.Options` Array(String) COMMENT 'Symbol options (comma separated list)',
+    `Groups.Names` Array(LowCardinality(String)) COMMENT 'Group name',
+    `Groups.Scores` Array(Float32) COMMENT 'Group score',
     ScanTimeReal UInt32 COMMENT 'Request time in milliseconds',
     ScanTimeVirtual UInt32 COMMENT 'Deprecated do not use',
     AuthUser String COMMENT 'Username for authenticated SMTP client',
@@ -226,6 +228,15 @@ local migrations = {
     -- New version
     [[INSERT INTO rspamd_version (Version) Values (7)]],
   },
+  [7] = {
+    -- Add new columns
+    [[ALTER TABLE rspamd
+      ADD COLUMN `Groups.Names` Array(LowCardinality(String)) AFTER `Symbols.Options`,
+      ADD COLUMN `Groups.Scores` Array(Float32) AFTER `Groups.Names`
+    ]],
+    -- New version
+    [[INSERT INTO rspamd_version (Version) Values (8)]],
+  },
 }
 
 local predefined_actions = {
@@ -314,6 +325,14 @@ local function clickhouse_symbols_row(res)
   for _,v in ipairs(fields) do table.insert(res, v) end
 end
 
+local function clickhouse_groups_row(res)
+  local fields = {
+    'Groups.Names',
+    'Groups.Scores',
+  }
+  for _,v in ipairs(fields) do table.insert(res, v) end
+end
+
 local function clickhouse_asn_row(res)
   local fields = {
     'ASN',
@@ -398,6 +417,7 @@ local function clickhouse_send_data(task, ev_base)
 
   if settings.enable_symbols then
     clickhouse_symbols_row(fields)
+    clickhouse_groups_row(fields)
   end
 
   send_data('generic data', data_rows,
@@ -787,6 +807,17 @@ local function clickhouse_collect(task)
     table.insert(row, syms_tab)
     table.insert(row, scores_tab)
     table.insert(row, options_tab)
+
+    -- Groups data
+    local groups = task:get_groups()
+    local groups_tab = {}
+    local gr_scores_tab = {}
+    for gr,sc in pairs(groups) do
+      table.insert(groups_tab, gr)
+      table.insert(gr_scores_tab, sc)
+    end
+    table.insert(row, groups_tab)
+    table.insert(row, gr_scores_tab)
   end
 
   -- Custom data
diff --git a/test/functional/cases/210_clickhouse/001_migration.robot b/test/functional/cases/210_clickhouse/001_migration.robot
index 5abda0881..8c51a3480 100644
--- a/test/functional/cases/210_clickhouse/001_migration.robot
+++ b/test/functional/cases/210_clickhouse/001_migration.robot
@@ -36,7 +36,9 @@ Migration
     # Added in schema version 7
     Column should exist    rspamd    Helo
     Column should exist    rspamd    SMTPRecipients
-    Schema version should be    7
+    # Added in schema version 8
+    Column should exist    rspamd    Groups.Scores
+    Schema version should be    8
 
 Retention
     Upload new schema        ${TESTDIR}/data/schema_2/schema.sql