plugins/mod_storage_sql2.lua

   1
   2 local json = require "util.json";
   3 local resolve_relative_path = require "core.configmanager".resolve_relative_path;
   4
   5 local unpack = unpack
   6 local function iterator(result)
   7         return function(result)
   8                 local row = result();
   9                 if row ~= nil then
  10                         return unpack(row);
  11                 end
  12         end, result, nil;
  13 end
  14
  15 local mod_sql = module:require("sql");
  16 local params = module:get_option("sql");
  17
  18 local engine; -- TODO create engine
  19
  20 local function create_table()
  21         --[[local Table,Column,Index = mod_sql.Table,mod_sql.Column,mod_sql.Index;
  22         local ProsodyTable = Table {
  23                 name="prosody";
  24                 Column { name="host", type="TEXT", nullable=false };
  25                 Column { name="user", type="TEXT", nullable=false };
  26                 Column { name="store", type="TEXT", nullable=false };
  27                 Column { name="key", type="TEXT", nullable=false };
  28                 Column { name="type", type="TEXT", nullable=false };
  29                 Column { name="value", type="TEXT", nullable=false };
  30                 Index { name="prosody_index", "host", "user", "store", "key" };
  31         };
  32         engine:transaction(function()
  33                 ProsodyTable:create(engine);
  34         end);]]
  35         if not module:get_option("sql_manage_tables", true) then
  36                 return;
  37         end
  38
  39         local create_sql = "CREATE TABLE `prosody` (`host` TEXT, `user` TEXT, `store` TEXT, `key` TEXT, `type` TEXT, `value` TEXT);";
  40         if params.driver == "PostgreSQL" then
  41                 create_sql = create_sql:gsub("`", "\"");
  42         elseif params.driver == "MySQL" then
  43                 create_sql = create_sql:gsub("`value` TEXT", "`value` MEDIUMTEXT")
  44                         :gsub(";$", " CHARACTER SET 'utf8' COLLATE 'utf8_bin';");
  45         end
  46
  47         local index_sql = "CREATE INDEX `prosody_index` ON `prosody` (`host`, `user`, `store`, `key`)";
  48         if params.driver == "PostgreSQL" then
  49                 index_sql = index_sql:gsub("`", "\"");
  50         elseif params.driver == "MySQL" then
  51                 index_sql = index_sql:gsub("`([,)])", "`(20)%1");
  52         end
  53
  54         local success,err = engine:transaction(function()
  55                 engine:execute(create_sql);
  56                 engine:execute(index_sql);
  57         end);
  58         if not success then -- so we failed to create
  59                 if params.driver == "MySQL" then
  60                         success,err = engine:transaction(function()
  61                                 local result = engine:execute("SHOW COLUMNS FROM prosody WHERE Field='value' and Type='text'");
  62                                 if result:rowcount() > 0 then
  63                                         module:log("info", "Upgrading database schema...");
  64                                         engine:execute("ALTER TABLE prosody MODIFY COLUMN `value` MEDIUMTEXT");
  65                                         module:log("info", "Database table automatically upgraded");
  66                                 end
  67                                 return true;
  68                         end);
  69                         if not success then
  70                                 module:log("error", "Failed to check/upgrade database schema (%s), please see "
  71                                         .."http://prosody.im/doc/mysql for help",
  72                                         err or "unknown error");
  73                         end
  74                 end
  75         end
  76 end
  77 local function set_encoding()
  78         if params.driver ~= "SQLite3" then
  79                 local set_names_query = "SET NAMES 'utf8';";
  80                 if params.driver == "MySQL" then
  81                         set_names_query = set_names_query:gsub(";$", " COLLATE 'utf8_bin';");
  82                 end
  83                 local success,err = engine:transaction(function() return engine:execute(set_names_query); end);
  84                 if not success then
  85                         module:log("error", "Failed to set database connection encoding to UTF8: %s", err);
  86                         return;
  87                 end
  88                 if params.driver == "MySQL" then
  89                         -- COMPAT w/pre-0.9: Upgrade tables to UTF-8 if not already
  90                         local check_encoding_query = "SELECT `COLUMN_NAME`,`COLUMN_TYPE` FROM `information_schema`.`columns` WHERE `TABLE_NAME`='prosody' AND ( `CHARACTER_SET_NAME`!='utf8' OR `COLLATION_NAME`!='utf8_bin' );";
  91                         local success,err = engine:transaction(function()
  92                                 local result = engine:execute(check_encoding_query);
  93                                 local n_bad_columns = result:rowcount();
  94                                 if n_bad_columns > 0 then
  95                                         module:log("warn", "Found %d columns in prosody table requiring encoding change, updating now...", n_bad_columns);
  96                                         local fix_column_query1 = "ALTER TABLE `prosody` CHANGE `%s` `%s` BLOB;";
  97                                         local fix_column_query2 = "ALTER TABLE `prosody` CHANGE `%s` `%s` %s CHARACTER SET 'utf8' COLLATE 'utf8_bin';";
  98                                         for row in result:rows() do
  99                                                 local column_name, column_type = unpack(row);
 100                                                 engine:execute(fix_column_query1:format(column_name, column_name));
 101                                                 engine:execute(fix_column_query2:format(column_name, column_name, column_type));
 102                                         end
 103                                         module:log("info", "Database encoding upgrade complete!");
 104                                 end
 105                         end);
 106                         local success,err = engine:transaction(function() return engine:execute(check_encoding_query); end);
 107                         if not success then
 108                                 module:log("error", "Failed to check/upgrade database encoding: %s", err or "unknown error");
 109                         end
 110                 end
 111         end
 112 end
 113
 114 do -- process options to get a db connection
 115         params = params or { driver = "SQLite3" };
 116
 117         if params.driver == "SQLite3" then
 118                 params.database = resolve_relative_path(prosody.paths.data or ".", params.database or "prosody.sqlite");
 119         end
 120
 121         assert(params.driver and params.database, "Both the SQL driver and the database need to be specified");
 122
 123         --local dburi = db2uri(params);
 124         engine = mod_sql:create_engine(params);
 125
 126         -- Encoding mess
 127         set_encoding();
 128
 129         -- Automatically create table, ignore failure (table probably already exists)
 130         create_table();
 131 end
 132
 133 local function serialize(value)
 134         local t = type(value);
 135         if t == "string" or t == "boolean" or t == "number" then
 136                 return t, tostring(value);
 137         elseif t == "table" then
 138                 local value,err = json.encode(value);
 139                 if value then return "json", value; end
 140                 return nil, err;
 141         end
 142         return nil, "Unhandled value type: "..t;
 143 end
 144 local function deserialize(t, value)
 145         if t == "string" then return value;
 146         elseif t == "boolean" then
 147                 if value == "true" then return true;
 148                 elseif value == "false" then return false; end
 149         elseif t == "number" then return tonumber(value);
 150         elseif t == "json" then
 151                 return json.decode(value);
 152         end
 153 end
 154
 155 local host = module.host;
 156 local user, store;
 157
 158 local function keyval_store_get()
 159         local haveany;
 160         local result = {};
 161         for row in engine:select("SELECT `key`,`type`,`value` FROM `prosody` WHERE `host`=? AND `user`=? AND `store`=?", host, user or "", store) do
 162                 haveany = true;
 163                 local k = row[1];
 164                 local v = deserialize(row[2], row[3]);
 165                 if k and v then
 166                         if k ~= "" then result[k] = v; elseif type(v) == "table" then
 167                                 for a,b in pairs(v) do
 168                                         result[a] = b;
 169                                 end
 170                         end
 171                 end
 172         end
 173         if haveany then
 174                 return result;
 175         end
 176 end
 177 local function keyval_store_set(data)
 178         engine:delete("DELETE FROM `prosody` WHERE `host`=? AND `user`=? AND `store`=?", host, user or "", store);
 179
 180         if data and next(data) ~= nil then
 181                 local extradata = {};
 182                 for key, value in pairs(data) do
 183                         if type(key) == "string" and key ~= "" then
 184                                 local t, value = serialize(value);
 185                                 assert(t, value);
 186                                 engine:insert("INSERT INTO `prosody` (`host`,`user`,`store`,`key`,`type`,`value`) VALUES (?,?,?,?,?,?)", host, user or "", store, key, t, value);
 187                         else
 188                                 extradata[key] = value;
 189                         end
 190                 end
 191                 if next(extradata) ~= nil then
 192                         local t, extradata = serialize(extradata);
 193                         assert(t, extradata);
 194                         engine:insert("INSERT INTO `prosody` (`host`,`user`,`store`,`key`,`type`,`value`) VALUES (?,?,?,?,?,?)", host, user or "", store, "", t, extradata);
 195                 end
 196         end
 197         return true;
 198 end
 199
 200 local keyval_store = {};
 201 keyval_store.__index = keyval_store;
 202 function keyval_store:get(username)
 203         user,store = username,self.store;
 204         return select(2, engine:transaction(keyval_store_get));
 205 end
 206 function keyval_store:set(username, data)
 207         user,store = username,self.store;
 208         return engine:transaction(function()
 209                 return keyval_store_set(data);
 210         end);
 211 end
 212 function keyval_store:users()
 213         local ok, result = engine:transaction(function()
 214                 return engine:select("SELECT DISTINCT `user` FROM `prosody` WHERE `host`=? AND `store`=?", host, self.store);
 215         end);
 216         if not ok then return ok, result end
 217         return iterator(result);
 218 end
 219
 220 local driver = {};
 221
 222 function driver:open(store, typ)
 223         if not typ then -- default key-value store
 224                 return setmetatable({ store = store }, keyval_store);
 225         end
 226         return nil, "unsupported-store";
 227 end
 228
 229 function driver:stores(username)
 230         local sql = "SELECT DISTINCT `store` FROM `prosody` WHERE `host`=? AND `user`" ..
 231                 (username == true and "!=?" or "=?");
 232         if username == true or not username then
 233                 username = "";
 234         end
 235         local ok, result = engine:transaction(function()
 236                 return engine:select(sql, host, username);
 237         end);
 238         if not ok then return ok, result end
 239         return iterator(result);
 240 end
 241
 242 function driver:purge(username)
 243         return engine:transaction(function()
 244                 local stmt,err = engine:delete("DELETE FROM `prosody` WHERE `host`=? AND `user`=?", host, username);
 245                 return true,err;
 246         end);
 247 end
 248
 249 module:provides("storage", driver);
 250
 251