From a79ffcd17eb2fcc8b5ee0780b75d5b74d8e2a24e Mon Sep 17 00:00:00 2001
From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Wed, 12 Jul 2017 18:29:09 +0100
Subject: [PATCH] [Feature] Automatically create tables in clickhouse

---
 src/plugins/lua/clickhouse.lua | 52 ++++++++++++++++++++++++++++------
 1 file changed, 44 insertions(+), 8 deletions(-)

diff --git a/src/plugins/lua/clickhouse.lua b/src/plugins/lua/clickhouse.lua
index eb4ad57b7..e424b9d67 100644
--- a/src/plugins/lua/clickhouse.lua
+++ b/src/plugins/lua/clickhouse.lua
@@ -57,10 +57,9 @@ local settings = {
   use_https = false,
 }
 
---[[
-Clickhouse schema:
-
-CREATE TABLE rspamd
+local clickhouse_schema = {
+rspamd = [[
+CREATE TABLE rspamd IF NOT EXISTS
 (
     Date Date,
     TS DateTime,
@@ -85,8 +84,10 @@ CREATE TABLE rspamd
     ListId String,
     Digest FixedString(32)
 ) ENGINE = MergeTree(Date, (TS, From), 8192)
+]],
 
-CREATE TABLE rspamd_attachments (
+  attachments = [[
+CREATE TABLE rspamd_attachments IF NOT EXISTS (
     Date Date,
     Digest FixedString(32),
     `Attachments.FileName` Array(String),
@@ -94,23 +95,29 @@ CREATE TABLE rspamd_attachments (
     `Attachments.Length` Array(UInt32),
     `Attachments.Digest` Array(FixedString(16))
 ) ENGINE = MergeTree(Date, Digest, 8192)
+]],
 
-CREATE TABLE rspamd_urls (
+  urls = [[
+CREATE TABLE rspamd_urls IF NOT EXISTS (
     Date Date,
     Digest FixedString(32),
     `Urls.Tld` Array(String),
     `Urls.Url` Array(String)
 ) ENGINE = MergeTree(Date, Digest, 8192)
+]],
 
-CREATE TABLE rspamd_asn (
+  asn = [[
+CREATE TABLE rspamd_asn IF NOT EXISTS (
     Date Date,
     Digest FixedString(32),
     ASN String,
     Country FixedString(2),
     IPNet String
 ) ENGINE = MergeTree(Date, Digest, 8192)
+]],
 
-CREATE TABLE rspamd_symbols (
+  symbols = [[
+CREATE TABLE rspamd_symbols IF NOT EXISTS (
     Date Date,
     Digest FixedString(32),
     `Symbols.Names` Array(String),
@@ -118,6 +125,7 @@ CREATE TABLE rspamd_symbols (
     `Symbols.Options` Array(String)
 ) ENGINE = MergeTree(Date, Digest, 8192)
 ]]
+}
 
 local function clickhouse_main_row(tname)
   local fields = {
@@ -656,5 +664,33 @@ if opts then
           clickhouse_send_data(task)
         end
       end)
+      -- Create tables on load
+      rspamd_config:add_on_load(function(cfg, ev_base, worker)
+        local function http_cb(err_message, code, _, _)
+          if code ~= 200 or err_message then
+            rspamd_logger.errx(task, "cannot create table in clickhouse server %s: %s:%s",
+              settings['server'], code, err_message)
+          end
+        end
+
+        local function send_req(elt, sql)
+          if not rspamd_http.request({
+            ev_base = ev_base,
+            config = cfg,
+            url = connect_prefix .. settings['server'],
+            body = sql,
+            callback = http_cb,
+            mime_type = 'text/plain',
+            timeout = settings['timeout'],
+          }) then
+            rspamd_logger.errx(task, "cannot create table %s in clickhouse server %s: cannot make request",
+              elt, settings['server'])
+          end
+        end
+
+        for tab,sql in pairs(clickhouse_schema) do
+          send_req(tab, sql)
+        end
+      end)
     end
 end
-- 
2.39.5