plugins/mod_storage_internal.lua

   1 local cache = require "util.cache";
   2 local datamanager = require "core.storagemanager".olddm;
   3 local array = require "util.array";
   4 local datetime = require "util.datetime";
   5 local st = require "util.stanza";
   6 local now = require "util.time".now;
   7 local id = require "util.id".medium;
   8 local jid_join = require "util.jid".join;
   9
  10 local host = module.host;
  11
  12 local archive_item_limit = module:get_option_number("storage_archive_item_limit", 10000);
  13 local archive_item_count_cache = cache.new(module:get_option("storage_archive_item_limit_cache_size", 1000));
  14
  15 local driver = {};
  16
  17 function driver:open(store, typ)
  18         local mt = self[typ or "keyval"]
  19         if not mt then
  20                 return nil, "unsupported-store";
  21         end
  22         return setmetatable({ store = store, type = typ }, mt);
  23 end
  24
  25 function driver:stores(username) -- luacheck: ignore 212/self
  26         return datamanager.stores(username, host);
  27 end
  28
  29 function driver:purge(user) -- luacheck: ignore 212/self
  30         return datamanager.purge(user, host);
  31 end
  32
  33 local keyval = { };
  34 driver.keyval = { __index = keyval };
  35
  36 function keyval:get(user)
  37         return datamanager.load(user, host, self.store);
  38 end
  39
  40 function keyval:set(user, data)
  41         return datamanager.store(user, host, self.store, data);
  42 end
  43
  44 function keyval:users()
  45         return datamanager.users(host, self.store, self.type);
  46 end
  47
  48 local archive = {};
  49 driver.archive = { __index = archive };
  50
  51 archive.caps = {
  52         total = true;
  53         quota = archive_item_limit;
  54         truncate = true;
  55 };
  56
  57 function archive:append(username, key, value, when, with)
  58         when = when or now();
  59         if not st.is_stanza(value) then
  60                 return nil, "unsupported-datatype";
  61         end
  62         value = st.preserialize(st.clone(value));
  63         value.when = when;
  64         value.with = with;
  65         value.attr.stamp = datetime.datetime(when);
  66         value.attr.stamp_legacy = datetime.legacy(when);
  67
  68         local cache_key = jid_join(username, host, self.store);
  69         local item_count = archive_item_count_cache:get(cache_key);
  70
  71         if key then
  72                 local items, err = datamanager.list_load(username, host, self.store);
  73                 if not items and err then return items, err; end
  74
  75                 -- Check the quota
  76                 item_count = items and #items or 0;
  77                 archive_item_count_cache:set(cache_key, item_count);
  78                 if item_count >= archive_item_limit then
  79                         module:log("debug", "%s reached or over quota, not adding to store", username);
  80                         return nil, "quota-limit";
  81                 end
  82
  83                 if items then
  84                         -- Filter out any item with the same key as the one being added
  85                         items = array(items);
  86                         items:filter(function (item)
  87                                 return item.key ~= key;
  88                         end);
  89
  90                         value.key = key;
  91                         items:push(value);
  92                         local ok, err = datamanager.list_store(username, host, self.store, items);
  93                         if not ok then return ok, err; end
  94                         archive_item_count_cache:set(cache_key, #items);
  95                         return key;
  96                 end
  97         else
  98                 if not item_count then -- Item count not cached?
  99                         -- We need to load the list to get the number of items currently stored
 100                         local items, err = datamanager.list_load(username, host, self.store);
 101                         if not items and err then return items, err; end
 102                         item_count = items and #items or 0;
 103                         archive_item_count_cache:set(cache_key, item_count);
 104                 end
 105                 if item_count >= archive_item_limit then
 106                         module:log("debug", "%s reached or over quota, not adding to store", username);
 107                         return nil, "quota-limit";
 108                 end
 109                 key = id();
 110         end
 111
 112         module:log("debug", "%s has %d items out of %d limit in store %s", username, item_count, archive_item_limit, self.store);
 113
 114         value.key = key;
 115
 116         local ok, err = datamanager.list_append(username, host, self.store, value);
 117         if not ok then return ok, err; end
 118         archive_item_count_cache:set(cache_key, item_count+1);
 119         return key;
 120 end
 121
 122 function archive:find(username, query)
 123         local items, err = datamanager.list_load(username, host, self.store);
 124         if not items then
 125                 if err then
 126                         return items, err;
 127                 elseif query then
 128                         if query.before or query.after then
 129                                 return nil, "item-not-found";
 130                         end
 131                         if query.total then
 132                                 return function () end, 0;
 133                         end
 134                 end
 135                 return function () end;
 136         end
 137         local count = nil;
 138         local i = 0;
 139         if query then
 140                 items = array(items);
 141                 if query.key then
 142                         items:filter(function (item)
 143                                 return item.key == query.key;
 144                         end);
 145                 end
 146                 if query.with then
 147                         items:filter(function (item)
 148                                 return item.with == query.with;
 149                         end);
 150                 end
 151                 if query.start then
 152                         items:filter(function (item)
 153                                 return item.when >= query.start;
 154                         end);
 155                 end
 156                 if query["end"] then
 157                         items:filter(function (item)
 158                                 return item.when <= query["end"];
 159                         end);
 160                 end
 161                 if query.total then
 162                         count = #items;
 163                 end
 164                 if query.reverse then
 165                         items:reverse();
 166                         if query.before then
 167                                 local found = false;
 168                                 for j = 1, #items do
 169                                         if (items[j].key or tostring(j)) == query.before then
 170                                                 found = true;
 171                                                 i = j;
 172                                                 break;
 173                                         end
 174                                 end
 175                                 if not found then
 176                                         return nil, "item-not-found";
 177                                 end
 178                         end
 179                 elseif query.after then
 180                         local found = false;
 181                         for j = 1, #items do
 182                                 if (items[j].key or tostring(j)) == query.after then
 183                                         found = true;
 184                                         i = j;
 185                                         break;
 186                                 end
 187                         end
 188                         if not found then
 189                                 return nil, "item-not-found";
 190                         end
 191                 end
 192                 if query.limit and #items - i > query.limit then
 193                         items[i+query.limit+1] = nil;
 194                 end
 195         end
 196         return function ()
 197                 i = i + 1;
 198                 local item = items[i];
 199                 if not item then return; end
 200                 local key = item.key or tostring(i);
 201                 local when = item.when or datetime.parse(item.attr.stamp);
 202                 local with = item.with;
 203                 item.key, item.when, item.with = nil, nil, nil;
 204                 item.attr.stamp = nil;
 205                 item.attr.stamp_legacy = nil;
 206                 item = st.deserialize(item);
 207                 return key, item, when, with;
 208         end, count;
 209 end
 210
 211 function archive:dates(username)
 212         local items, err = datamanager.list_load(username, host, self.store);
 213         if not items then return items, err; end
 214         return array(items):pluck("when"):map(datetime.date):unique();
 215 end
 216
 217 function archive:summary(username, query)
 218         local iter, err = self:find(username, query)
 219         if not iter then return iter, err; end
 220         local counts = {};
 221         local earliest = {};
 222         local latest = {};
 223         local body = {};
 224         for _, stanza, when, with in iter do
 225                 counts[with] = (counts[with] or 0) + 1;
 226                 if earliest[with] == nil then
 227                         earliest[with] = when;
 228                 end
 229                 latest[with] = when;
 230                 body[with] = stanza:get_child_text("body") or body[with];
 231         end
 232         return {
 233                 counts = counts;
 234                 earliest = earliest;
 235                 latest = latest;
 236                 body = body;
 237         };
 238 end
 239
 240 function archive:users()
 241         return datamanager.users(host, self.store, "list");
 242 end
 243
 244 function archive:delete(username, query)
 245         local cache_key = jid_join(username, host, self.store);
 246         if not query or next(query) == nil then
 247                 archive_item_count_cache:set(cache_key, nil);
 248                 return datamanager.list_store(username, host, self.store, nil);
 249         end
 250         local items, err = datamanager.list_load(username, host, self.store);
 251         if not items then
 252                 if err then
 253                         return items, err;
 254                 end
 255                 archive_item_count_cache:set(cache_key, 0);
 256                 -- Store is empty
 257                 return 0;
 258         end
 259         items = array(items);
 260         local count_before = #items;
 261         if query then
 262                 if query.key then
 263                         items:filter(function (item)
 264                                 return item.key ~= query.key;
 265                         end);
 266                 end
 267                 if query.with then
 268                         items:filter(function (item)
 269                                 return item.with ~= query.with;
 270                         end);
 271                 end
 272                 if query.start then
 273                         items:filter(function (item)
 274                                 return item.when < query.start;
 275                         end);
 276                 end
 277                 if query["end"] then
 278                         items:filter(function (item)
 279                                 return item.when > query["end"];
 280                         end);
 281                 end
 282                 if query.truncate and #items > query.truncate then
 283                         if query.reverse then
 284                                 -- Before: { 1, 2, 3, 4, 5, }
 285                                 -- After: { 1, 2, 3 }
 286                                 for i = #items, query.truncate + 1, -1 do
 287                                         items[i] = nil;
 288                                 end
 289                         else
 290                                 -- Before: { 1, 2, 3, 4, 5, }
 291                                 -- After: { 3, 4, 5 }
 292                                 local offset = #items - query.truncate;
 293                                 for i = 1, #items do
 294                                         items[i] = items[i+offset];
 295                                 end
 296                         end
 297                 end
 298         end
 299         local count = count_before - #items;
 300         if count == 0 then
 301                 return 0; -- No changes, skip write
 302         end
 303         local ok, err = datamanager.list_store(username, host, self.store, items);
 304         if not ok then return ok, err; end
 305         archive_item_count_cache:set(cache_key, #items);
 306         return count;
 307 end
 308
 309 module:provides("storage", driver);