mod_storage_gdbm/mod_storage_gdbm.lua

   1 -- mod_storage_gdbm
   2 -- Copyright (C) 2014-2015 Kim Alvefur
   3 --
   4 -- This file is MIT/X11 licensed.
   5 --
   6 -- Depends on lgdbm:
   7 -- http://webserver2.tecgraf.puc-rio.br/~lhf/ftp/lua/#lgdbm
   8 --
   9 -- luacheck: globals open purge
  10
  11 local gdbm = require"gdbm";
  12 local path = require"util.paths";
  13 local lfs = require"lfs";
  14 local st = require"util.stanza";
  15 local uuid = require"util.uuid".generate;
  16
  17 local serialization = require"util.serialization";
  18 local serialize = serialization.serialize;
  19 local deserialize = serialization.deserialize;
  20
  21 local g_set, g_get, g_del = gdbm.replace, gdbm.fetch, gdbm.delete;
  22 local g_first, g_next = gdbm.firstkey, gdbm.nextkey;
  23
  24 local empty = {};
  25
  26 local function id(v) return v; end
  27
  28 local function is_stanza(s)
  29         return getmetatable(s) == st.stanza_mt;
  30 end
  31
  32 local base_path = path.resolve_relative_path(prosody.paths.data, module.host);
  33 lfs.mkdir(base_path);
  34
  35 local cache = {};
  36
  37 local keyval = {};
  38 local keyval_mt = { __index = keyval, suffix = ".db" };
  39
  40 function keyval:set(user, value)
  41         if type(value) == "table" and next(value) == nil then
  42                 value = nil;
  43         end
  44         if value ~= nil then
  45                 value = serialize(value);
  46         end
  47         local ok, err = (value and g_set or g_del)(self._db, user or "@", value);
  48         if not ok then return nil, err; end
  49         return true;
  50 end
  51
  52 function keyval:get(user)
  53         local data, err = g_get(self._db, user or "@");
  54         if not data then return nil, err; end
  55         return deserialize(data);
  56 end
  57
  58 local function g_keys(db, key)
  59         return (key == nil and g_first or g_next)(db, key);
  60 end
  61
  62 function keyval:users()
  63         return g_keys, self._db, nil;
  64 end
  65
  66 local archive = {};
  67 local archive_mt = { __index = archive, suffix = ".adb" };
  68
  69 archive.get = keyval.get;
  70 archive.set = keyval.set;
  71
  72 function archive:append(username, key, value, when, with)
  73         if type(when) ~= "number" then
  74                 when, with, value = value, when, with;
  75         end
  76
  77         key = key or uuid();
  78         local meta = self:get(username);
  79         if not meta then
  80                 meta = {};
  81         end
  82         local i = meta[key] or #meta+1;
  83         local type;
  84         if is_stanza(value) then
  85                 type, value = "stanza", st.preserialize(value);
  86         end
  87         meta[i] = { key = key, when = when, with = with, type = type };
  88         meta[key] = i;
  89         local prefix = (username or "@") .. "#";
  90         local ok, err = self:set(prefix..key, value);
  91         if not ok then return nil, err; end
  92         ok, err = self:set(username, meta);
  93         if not ok then return nil, err; end
  94         return key;
  95 end
  96
  97 local deserialize_map = {
  98         stanza = st.deserialize;
  99 };
 100
 101 function archive:find(username, query)
 102         query = query or empty;
 103         local meta = self:get(username) or empty;
 104         local prefix = (username or "@") .. "#";
 105         local reverse = query.reverse;
 106         local step = reverse and -1 or 1;
 107         local first = meta[query.after];
 108         local last = meta[query.before];
 109         -- we don't want to include the actual 'after' or 'before' item
 110         if first then first = first + 1; end
 111         if last then last = last - 1; end
 112         first, last = first or 1, last or #meta;
 113         if reverse then first, last = last, first; end
 114         local limit = query.limit;
 115         local count = 0;
 116         return function ()
 117                 if limit and count >= limit then return end
 118                 local item, value;
 119                 for i = first, last, step do
 120                         item = meta[i];
 121                         if (not query.with or item.with == query.with)
 122                         and (not query.start or item.when >= query.start)
 123                         and (not query["end"] or item.when <= query["end"]) then
 124                                 first = i + step; count = count + 1;
 125                                 value = self:get(prefix..item.key);
 126                                 return item.key, (deserialize_map[item.type] or id)(value), item.when, item.with;
 127                         end
 128                 end
 129         end
 130 end
 131
 132 local drivers = {
 133         keyval = keyval_mt;
 134         archive = archive_mt;
 135 }
 136
 137 function open(_, store, typ)
 138         typ = typ or "keyval";
 139         local driver_mt = drivers[typ];
 140         if not driver_mt then
 141                 return nil, "unsupported-store";
 142         end
 143
 144         local db_path = path.join(base_path, store) .. driver_mt.suffix;
 145
 146         local db = cache[db_path];
 147         if not db then
 148                 db = assert(gdbm.open(db_path, "c"));
 149                 cache[db_path] = db;
 150         end
 151         return setmetatable({ _db = db; _path = db_path; store = store, type = typ }, driver_mt);
 152 end
 153
 154 function purge(_, user)
 155         for dir in lfs.dir(base_path) do
 156                 local name, ext = dir:match("^(.-)%.a?db$");
 157                 if ext == ".db" then
 158                         open(_, name, "keyval"):set(user, nil);
 159                 elseif ext == ".adb" then
 160                         open(_, name, "archive"):delete(user);
 161                 end
 162         end
 163         return true;
 164 end
 165
 166 function module.unload()
 167         for db_path, db in pairs(cache) do
 168                 module:log("debug", "Closing db at %q", db_path);
 169                 gdbm.reorganize(db);
 170                 gdbm.sync(db);
 171                 gdbm.close(db);
 172         end
 173 end
 174
 175 module:provides"storage";