Diff

--- a/plugins/mod_storage_internal.lua	Thu Oct 01 15:08:58 2020 +0100
+++ b/plugins/mod_storage_internal.lua	Wed Sep 30 09:50:33 2020 +0100
@@ -1,12 +1,17 @@
+local cache = require "util.cache";
 local datamanager = require "core.storagemanager".olddm;
 local array = require "util.array";
 local datetime = require "util.datetime";
 local st = require "util.stanza";
 local now = require "util.time".now;
 local id = require "util.id".medium;
+local jid_join = require "util.jid".join;
 
 local host = module.host;
 
+local archive_item_limit = module:get_option_number("storage_archive_item_limit", 10000);
+local archive_item_count_cache = cache.new(module:get_option("storage_archive_item_limit_cache_size", 1000));
+
 local driver = {};
 
 function driver:open(store, typ)
@@ -43,6 +48,12 @@
 local archive = {};
 driver.archive = { __index = archive };
 
+archive.caps = {
+	total = true;
+	quota = archive_item_limit;
+	truncate = true;
+};
+
 function archive:append(username, key, value, when, with)
 	when = when or now();
 	if not st.is_stanza(value) then
@@ -54,28 +65,57 @@
 	value.attr.stamp = datetime.datetime(when);
 	value.attr.stamp_legacy = datetime.legacy(when);
 
+	local cache_key = jid_join(username, host, self.store);
+	local item_count = archive_item_count_cache:get(cache_key);
+
 	if key then
 		local items, err = datamanager.list_load(username, host, self.store);
 		if not items and err then return items, err; end
+
+		-- Check the quota
+		item_count = items and #items or 0;
+		archive_item_count_cache:set(cache_key, item_count);
+		if item_count >= archive_item_limit then
+			module:log("debug", "%s reached or over quota, not adding to store", username);
+			return nil, "quota-limit";
+		end
+
 		if items then
+			-- Filter out any item with the same key as the one being added
 			items = array(items);
 			items:filter(function (item)
 				return item.key ~= key;
 			end);
+
 			value.key = key;
 			items:push(value);
 			local ok, err = datamanager.list_store(username, host, self.store, items);
 			if not ok then return ok, err; end
+			archive_item_count_cache:set(cache_key, #items);
 			return key;
 		end
 	else
+		if not item_count then -- Item count not cached?
+			-- We need to load the list to get the number of items currently stored
+			local items, err = datamanager.list_load(username, host, self.store);
+			if not items and err then return items, err; end
+			item_count = items and #items or 0;
+			archive_item_count_cache:set(cache_key, item_count);
+		end
+		if item_count >= archive_item_limit then
+			module:log("debug", "%s reached or over quota, not adding to store", username);
+			return nil, "quota-limit";
+		end
 		key = id();
 	end
 
+	module:log("debug", "%s has %d items out of %d limit in store %s", username, item_count, archive_item_limit, self.store);
+
 	value.key = key;
 
 	local ok, err = datamanager.list_append(username, host, self.store, value);
 	if not ok then return ok, err; end
+	archive_item_count_cache:set(cache_key, item_count+1);
 	return key;
 end
 
@@ -84,12 +124,18 @@
 	if not items then
 		if err then
 			return items, err;
-		else
-			return function () end, 0;
+		elseif query then
+			if query.before or query.after then
+				return nil, "item-not-found";
+			end
+			if query.total then
+				return function () end, 0;
+			end
 		end
+		return function () end;
 	end
-	local count = #items;
-	local i = 0;
+	local count = nil;
+	local i, last_key = 0;
 	if query then
 		items = array(items);
 		if query.key then
@@ -114,24 +160,38 @@
 				return when <= query["end"];
 			end);
 		end
-		count = #items;
+		if query.total then
+			count = #items;
+		end
 		if query.reverse then
 			items:reverse();
 			if query.before then
-				for j = 1, count do
+				local found = false;
+				for j = 1, #items do
 					if (items[j].key or tostring(j)) == query.before then
+						found = true;
 						i = j;
 						break;
 					end
 				end
+				if not found then
+					return nil, "item-not-found";
+				end
 			end
+		elseif query.before then
+			last_key = query.before;
 		elseif query.after then
-			for j = 1, count do
+			local found = false;
+			for j = 1, #items do
 				if (items[j].key or tostring(j)) == query.after then
+					found = true;
 					i = j;
 					break;
 				end
 			end
+			if not found then
+				return nil, "item-not-found";
+			end
 		end
 		if query.limit and #items - i > query.limit then
 			items[i+query.limit+1] = nil;
@@ -140,7 +200,9 @@
 	return function ()
 		i = i + 1;
 		local item = items[i];
-		if not item then return; end
+		if not item or (last_key and item.key == last_key) then
+			return;
+		end
 		local key = item.key or tostring(i);
 		local when = item.when or datetime.parse(item.attr.stamp);
 		local with = item.with;
@@ -152,14 +214,83 @@
 	end, count;
 end
 
+function archive:get(username, wanted_key)
+	local iter, err = self:find(username, { key = wanted_key })
+	if not iter then return iter, err; end
+	for key, stanza, when, with in iter do
+		if key == wanted_key then
+			return stanza, when, with;
+		end
+	end
+	return nil, "item-not-found";
+end
+
+function archive:set(username, key, new_value, new_when, new_with)
+	local items, err = datamanager.list_load(username, host, self.store);
+	if not items then
+		if err then
+			return items, err;
+		else
+			return nil, "item-not-found";
+		end
+	end
+
+	for i = 1, #items do
+		local old_item = items[i];
+		if old_item.key == key then
+			local item = st.preserialize(st.clone(new_value));
+
+			local when = new_when or old_item.when or datetime.parse(old_item.attr.stamp);
+			item.key = key;
+			item.when = when;
+			item.with = new_with or old_item.with;
+			item.attr.stamp = datetime.datetime(when);
+			item.attr.stamp_legacy = datetime.legacy(when);
+			items[i] = item;
+			return datamanager.list_store(username, host, self.store, items);
+		end
+	end
+
+	return nil, "item-not-found";
+end
+
 function archive:dates(username)
 	local items, err = datamanager.list_load(username, host, self.store);
 	if not items then return items, err; end
 	return array(items):pluck("when"):map(datetime.date):unique();
 end
 
+function archive:summary(username, query)
+	local iter, err = self:find(username, query)
+	if not iter then return iter, err; end
+	local counts = {};
+	local earliest = {};
+	local latest = {};
+	local body = {};
+	for _, stanza, when, with in iter do
+		counts[with] = (counts[with] or 0) + 1;
+		if earliest[with] == nil then
+			earliest[with] = when;
+		end
+		latest[with] = when;
+		body[with] = stanza:get_child_text("body") or body[with];
+	end
+	return {
+		counts = counts;
+		earliest = earliest;
+		latest = latest;
+		body = body;
+	};
+end
+
+function archive:users()
+	return datamanager.users(host, self.store, "list");
+end
+
 function archive:delete(username, query)
+	local cache_key = jid_join(username, host, self.store);
 	if not query or next(query) == nil then
+		archive_item_count_cache:set(cache_key, nil);
 		return datamanager.list_store(username, host, self.store, nil);
 	end
 	local items, err = datamanager.list_load(username, host, self.store);
@@ -167,6 +298,7 @@
 		if err then
 			return items, err;
 		end
+		archive_item_count_cache:set(cache_key, 0);
 		-- Store is empty
 		return 0;
 	end
@@ -216,6 +348,7 @@
 	end
 	local ok, err = datamanager.list_store(username, host, self.store, items);
 	if not ok then return ok, err; end
+	archive_item_count_cache:set(cache_key, #items);
 	return count;
 end