mod_log_messages_sql/mod_log_messages_sql.lua

   1 -- Based on mod_mam_sql
   2 -- Copyright (C) 2011-2012 Kim Alvefur
   3 --
   4 -- This file is MIT/X11 licensed.
   5
   6 local st = require "util.stanza";
   7 local jid_bare = require "util.jid".bare;
   8 local jid_split = require "util.jid".split;
   9
  10 local serialize = require"util.json".encode, require"util.json".decode;
  11 local tostring = tostring;
  12 local time_now = os.time;
  13
  14
  15 local table_name = module:get_option("message_log_sql_table", pcall(require, "util.cache") and "messages" or "prosodyarchive");
  16
  17 local sql, setsql, getsql = {};
  18 do -- SQL stuff
  19 local connection;
  20 local resolve_relative_path = require "core.configmanager".resolve_relative_path;
  21 local params = module:get_option("message_log_sql", module:get_option("sql"));
  22
  23
  24 local function test_connection()
  25         if not connection then return nil; end
  26         if connection:ping() then
  27                 return true;
  28         else
  29                 module:log("debug", "Database connection closed");
  30                 connection = nil;
  31         end
  32 end
  33 local function connect()
  34         if not test_connection() then
  35                 prosody.unlock_globals();
  36                 local dbh, err = DBI.Connect(
  37                         params.driver, params.database,
  38                         params.username, params.password,
  39                         params.host, params.port
  40                 );
  41                 prosody.lock_globals();
  42                 if not dbh then
  43                         module:log("debug", "Database connection failed: %s", tostring(err));
  44                         return nil, err;
  45                 end
  46                 module:log("debug", "Successfully connected to database");
  47                 dbh:autocommit(false); -- don't commit automatically
  48                 connection = dbh;
  49
  50         end
  51         return connection;
  52 end
  53
  54
  55 do -- process options to get a db connection
  56         local ok;
  57         prosody.unlock_globals();
  58         ok, DBI = pcall(require, "DBI");
  59         if not ok then
  60                 package.loaded["DBI"] = {};
  61                 module:log("error", "Failed to load the LuaDBI library for accessing SQL databases: %s", DBI);
  62                 module:log("error", "More information on installing LuaDBI can be found at http://prosody.im/doc/depends#luadbi");
  63         end
  64         prosody.lock_globals();
  65         if not ok or not DBI.Connect then
  66                 return; -- Halt loading of this module
  67         end
  68
  69         params = params or { driver = "SQLite3" };
  70
  71         if params.driver == "SQLite3" then
  72                 params.database = resolve_relative_path(prosody.paths.data or ".", params.database or "prosody.sqlite");
  73         end
  74
  75         assert(params.driver and params.database, "Both the SQL driver and the database need to be specified");
  76
  77         assert(connect());
  78
  79 end
  80
  81 function getsql(sql, ...)
  82         if params.driver == "PostgreSQL" then
  83                 sql = sql:gsub("`", "\"");
  84         end
  85         -- do prepared statement stuff
  86         local stmt, err = connection:prepare(sql);
  87         if not stmt and not test_connection() then error("connection failed"); end
  88         if not stmt then module:log("error", "QUERY FAILED: %s %s", err, debug.traceback()); return nil, err; end
  89         -- run query
  90         local ok, err = stmt:execute(...);
  91         if not ok and not test_connection() then error("connection failed"); end
  92         if not ok then return nil, err; end
  93
  94         return stmt;
  95 end
  96 function setsql(sql, ...)
  97         local stmt, err = getsql(sql, ...);
  98         if not stmt then return stmt, err; end
  99         return stmt:affected();
 100 end
 101 function sql.rollback(...)
 102         if connection then connection:rollback(); end -- FIXME check for rollback error?
 103         return ...;
 104 end
 105 function sql.commit(...)
 106         local ok, err = connection:commit();
 107         if not ok then
 108                 module:log("error", "SQL commit failed: %s", tostring(err));
 109                 return nil, "SQL commit failed: "..tostring(err);
 110         end
 111         return ...;
 112 end
 113
 114 end
 115
 116 -- Handle messages
 117 local function message_handler(event, c2s)
 118         local origin, stanza = event.origin, event.stanza;
 119         local orig_type = stanza.attr.type or "normal";
 120         local orig_to = stanza.attr.to;
 121         local orig_from = stanza.attr.from;
 122
 123         if not orig_from and c2s then
 124                 orig_from = origin.full_jid;
 125         end
 126         orig_to = orig_to or orig_from; -- Weird corner cases
 127
 128         -- Don't store messages of these types
 129         if orig_type == "error"
 130         or orig_type == "headline"
 131         or orig_type == "groupchat"
 132         or not stanza:get_child("body") then
 133                 return;
 134                 -- TODO Maybe headlines should be configurable?
 135         end
 136
 137         local store_user, store_host = jid_split(c2s and orig_from or orig_to);
 138         local target_jid = c2s and orig_to or orig_from;
 139         local target_bare = jid_bare(target_jid);
 140         local _, _, target_resource = jid_split(target_jid);
 141
 142         --local id = uuid();
 143         local when = time_now();
 144         -- And stash it
 145         local ok, err = setsql([[
 146         INSERT INTO `]]..table_name..[[`
 147         (`host`, `user`, `store`, `when`, `with`, `resource`, `stanza`)
 148         VALUES (?, ?, ?, ?, ?, ?, ?);
 149         ]], store_host, store_user, "message_log", when, target_bare, target_resource, serialize(st.preserialize(stanza)))
 150         if ok then
 151                 sql.commit();
 152         else
 153                 module:log("error", "SQL error: %s", err);
 154                 sql.rollback();
 155         end
 156 end
 157
 158 local function c2s_message_handler(event)
 159         return message_handler(event, true);
 160 end
 161
 162 -- Stanzas sent by local clients
 163 module:hook("pre-message/bare", c2s_message_handler, 2);
 164 module:hook("pre-message/full", c2s_message_handler, 2);
 165 -- Stanszas to local clients
 166 module:hook("message/bare", message_handler, 2);
 167 module:hook("message/full", message_handler, 2);
 168
 169 -- In the telnet console, run:
 170 -- >hosts["this host"].modules.mam_sql.environment.create_sql()
 171 function create_sql()
 172         local stm = assert(getsql([[
 173         CREATE TABLE `]]..table_name..[[` (
 174                 `host` TEXT,
 175                 `user` TEXT,
 176                 `store` TEXT,
 177                 `id` INTEGER PRIMARY KEY AUTOINCREMENT,
 178                 `when` INTEGER,
 179                 `with` TEXT,
 180                 `resource` TEXT,
 181                 `stanza` TEXT
 182         );
 183         CREATE INDEX `hus` ON `]]..table_name..[[` (`host`, `user`, `store`);
 184         CREATE INDEX `with` ON `]]..table_name..[[` (`with`);
 185         CREATE INDEX `thetime` ON `]]..table_name..[[` (`when`);
 186         ]]));
 187         stm:execute();
 188         sql.commit();
 189 end