plugins/mod_storage_sql_ejabberd.lua

   1
   2 local setmetatable = setmetatable;
   3 local error = error;
   4 local unpack = unpack;
   5 local module = module;
   6 local tostring = tostring;
   7 local pairs, next = pairs, next;
   8 local prosody = prosody;
   9 local assert = assert;
  10 local require = require;
  11 local st = require "util.stanza";
  12 local DBI = require "DBI";
  13
  14 -- connect to db
  15 local params = module:get_option("sql_ejabberd") or error("No sql_ejabberd config option");
  16 local database;
  17 do
  18         module:log("debug", "Opening database: %s", "dbi:"..params.driver..":"..params.database);
  19         prosody.unlock_globals();
  20         local dbh, err = DBI.Connect(
  21                 params.driver, params.database,
  22                 params.username, params.password,
  23                 params.host, params.port
  24         );
  25         prosody.lock_globals();
  26         assert(dbh, err);
  27         dbh:autocommit(true);
  28         database = dbh;
  29 end
  30
  31 -- initialize db
  32 local ejabberd_init = module:require("ejabberd_init");
  33 ejabberd_init.init(database);
  34
  35 local sqlcache = {};
  36 local function prepare(sql)
  37         module:log("debug", "query: %s", sql);
  38         local err;
  39         local r = sqlcache[sql];
  40         if not r then
  41                 r, err = database:prepare(sql);
  42                 if not r then error("Unable to prepare SQL statement: "..err); end
  43                 sqlcache[sql] = r;
  44         end
  45         return r;
  46 end
  47
  48 local _parse_xml = module:require("xmlparse");
  49 local function parse_xml(str)
  50         local s = _parse_xml(str);
  51         if s and not s.gsub then
  52                 return st.preserialize(s);
  53         end
  54 end
  55 local function unparse_xml(s)
  56         return tostring(st.deserialize(s));
  57 end
  58
  59
  60 local handlers = {};
  61
  62 handlers.accounts = {
  63         get = function(self, user)
  64                 local select = self:query("select password from users where username=? and host=?", user, self.host);
  65                 local row = select and select:fetch();
  66                 if row then return { password = row[1] }; end
  67         end;
  68         set = function(self, user, data)
  69                 if data and data.password then
  70                         return self:modify("update users set password=? where username=? and host=?", data.password, user, self.host)
  71                                 or self:modify("insert into users (username, host, password) values (?, ?, ?)", user, self.host, data.password);
  72                 else
  73                         return self:modify("delete from users where username=? and host=?", user, self.host);
  74                 end
  75         end;
  76 };
  77 handlers.vcard = {
  78         get = function(self, user)
  79                 local select = self:query("select vcard from vcard where username=? and host=?", user, self.host);
  80                 local row = select and select:fetch();
  81                 if row then return parse_xml(row[1]); end
  82         end;
  83         set = function(self, user, data)
  84                 if data then
  85                         data = unparse_xml(data);
  86                         return self:modify("update vcard set vcard=? where username=? and host=?", data, user, self.host)
  87                                 or self:modify("insert into vcard (username, host, vcard) values (?, ?, ?)", user, self.host, data);
  88                 else
  89                         return self:modify("delete from vcard where username=? and host=?", user, self.host);
  90                 end
  91         end;
  92 };
  93 handlers.private = {
  94         get = function(self, user)
  95                 local select = self:query("select namespace,data from private_storage where username=? and host=?", user, self.host);
  96                 if select then
  97                         local data = {};
  98                         for row in select:rows() do
  99                                 data[row[1]] = parse_xml(row[2]);
 100                         end
 101                         return data;
 102                 end
 103         end;
 104         set = function(self, user, data)
 105                 if data then
 106                         self:modify("delete from private_storage where username=? and host=?", user, self.host);
 107                         for namespace,text in pairs(data) do
 108                                 self:modify("insert into private_storage (username, host, namespace, data) values (?, ?, ?, ?)", user, self.host, namespace, unparse_xml(text));
 109                         end
 110                         return true;
 111                 else
 112                         return self:modify("delete from private_storage where username=? and host=?", user, self.host);
 113                 end
 114         end;
 115         -- TODO map_set, map_get
 116 };
 117 local subscription_map = { N = "none", B = "both", F = "from", T = "to" };
 118 local subscription_map_reverse = { none = "N", both = "B", from = "F", to = "T" };
 119 handlers.roster = {
 120         get = function(self, user)
 121                 local select = self:query("select jid,nick,subscription,ask,server,subscribe,type from rosterusers where username=?", user);
 122                 if select then
 123                         local roster = { pending = {} };
 124                         for row in select:rows() do
 125                                 local jid,nick,subscription,ask,server,subscribe,typ = unpack(row);
 126                                 local item = { groups = {} };
 127                                 if nick == "" then nick = nil; end
 128                                 item.nick = nick;
 129                                 item.subscription = subscription_map[subscription];
 130                                 if ask == "N" then ask = nil;
 131                                 elseif ask == "O" then ask = "subscribe"
 132                                 elseif ask == "I" then roster.pending[jid] = true; ask = nil;
 133                                 elseif ask == "B" then roster.pending[jid] = true; ask = "subscribe";
 134                                 else module:log("debug", "bad roster_item.ask: %s", ask); ask = nil; end
 135                                 item.ask = ask;
 136                                 roster[jid] = item;
 137                         end
 138
 139                         select = self:query("select jid,grp from rostergroups where username=?", user);
 140                         if select then
 141                                 for row in select:rows() do
 142                                         local jid,grp = unpack(row);
 143                                         if roster[jid] then roster[jid].groups[grp] = true; end
 144                                 end
 145                         end
 146                         select = self:query("select version from roster_version where username=?", user);
 147                         local row = select and select:fetch();
 148                         if row then
 149                                 roster[false] = { version = row[1]; };
 150                         end
 151                         return roster;
 152                 end
 153         end;
 154         set = function(self, user, data)
 155                 if data and next(data) ~= nil then
 156                         self:modify("delete from rosterusers where username=?", user);
 157                         self:modify("delete from rostergroups where username=?", user);
 158                         self:modify("delete from roster_version where username=?", user);
 159                         local done = {};
 160                         local pending = data.pending or {};
 161                         for jid,item in pairs(data) do
 162                                 if jid and jid ~= "pending" then
 163                                         local subscription = subscription_map_reverse[item.subscription];
 164                                         local ask;
 165                                         if pending[jid] then
 166                                                 if item.ask then ask = "B"; else ask = "I"; end
 167                                         else
 168                                                 if item.ask then ask = "O"; else ask = "N"; end
 169                                         end
 170                                         local r = self:modify("insert into rosterusers (username,jid,nick,subscription,ask,askmessage,server,subscribe) values (?, ?, ?, ?, ?, '', '', '')", user, jid, item.nick or "", subscription, ask);
 171                                         if not r then module:log("debug", "--- :( %s", tostring(r)); end
 172                                         done[jid] = true;
 173                                         for group in pairs(item.groups) do
 174                                                 self:modify("insert into rostergroups (username,jid,grp) values (?, ?, ?)", user, jid, group);
 175                                         end
 176                                 end
 177                         end
 178                         for jid in pairs(pending) do
 179                                 if not done[jid] then
 180                                         self:modify("insert into rosterusers (username,jid,nick,subscription,ask,askmessage,server,subscribe) values (?, ?, ?, ?, ?. ''. ''. '')", user, jid, "", "N", "I");
 181                                 end
 182                         end
 183                         local version = data[false] and data[false].version;
 184                         if version then
 185                                 self:modify("insert into roster_version (username,version) values (?, ?)", user, version);
 186                         end
 187                         return true;
 188                 else
 189                         self:modify("delete from rosterusers where username=?", user);
 190                         self:modify("delete from rostergroups where username=?", user);
 191                         self:modify("delete from roster_version where username=?", user);
 192                 end
 193         end;
 194 };
 195
 196 -----------------------------
 197 local driver = {};
 198 driver.__index = driver;
 199
 200 function driver:query(sql, ...)
 201         local stmt,err = prepare(sql);
 202         if not stmt then
 203                 module:log("error", "Failed to prepare SQL [[%s]], error: %s", sql, err);
 204                 return nil, err;
 205         end
 206         local ok, err = stmt:execute(...);
 207         if not ok then
 208                 module:log("error", "Failed to execute SQL [[%s]], error: %s", sql, err);
 209                 return nil, err;
 210         end
 211         return stmt;
 212 end
 213 function driver:modify(sql, ...)
 214         local stmt, err = self:query(sql, ...);
 215         if stmt and stmt:affected() > 0 then return stmt; end
 216         return nil, err;
 217 end
 218
 219 function driver:open(datastore, typ)
 220         local instance = setmetatable({ host = module.host, datastore = datastore }, self);
 221         local handler = handlers[datastore];
 222         if not handler then return nil; end
 223         for key,val in pairs(handler) do
 224                 instance[key] = val;
 225         end
 226         if instance.init then instance:init(); end
 227         return instance;
 228 end
 229
 230 -----------------------------
 231
 232 module:add_item("data-driver", driver);