Add commands for cleaning up mam archive

3 files changed, 88 insertions, 4 deletions

diff --git a/src/mod_mam.erl b/src/mod_mam.erl
index ac88c82c..7912be30 100644
--- a/src/mod_mam.erl
+++ b/src/mod_mam.erl
@@ -40,7 +40,7 @@
 	 muc_process_iq/2, muc_filter_message/3, message_is_archived/3,
 	 delete_old_messages/2, get_commands_spec/0, msg_to_el/4,
 	 get_room_config/4, set_room_option/3, offline_message/1, export/1,
-	 mod_options/1]).
+	 mod_options/1, remove_mam_for_user_with_peer/3, remove_mam_for_user/2]).
 
 -include("xmpp.hrl").
 -include("logger.hrl").
@@ -69,6 +69,7 @@
     {[{binary(), non_neg_integer(), xmlel()}], boolean(), non_neg_integer()}.
 -callback use_cache(binary()) -> boolean().
 -callback cache_nodes(binary()) -> [node()].
+-callback remove_from_archive(binary(), binary(), jid() | none) -> ok | {error, any()}.
 
 -optional_callbacks([use_cache/1, cache_nodes/1]).
 
@@ -258,6 +259,37 @@ remove_room(LServer, Name, Host) ->
     Mod:remove_room(LServer, LName, LHost),
     ok.
 
+remove_mam_for_user(User, Server) ->
+    LUser = jid:nodeprep(User),
+    LServer = jid:nameprep(Server),
+    Mod = gen_mod:db_mod(LServer, ?MODULE),
+    case Mod:remove_from_archive(LUser, LServer, none) of
+	ok ->
+	    {ok, <<"MAM archive removed">>};
+	{error, Bin} when is_binary(Bin) ->
+	    {error, Bin};
+	{error, _} ->
+	    {error, <<"Db returned error">>}
+    end.
+
+remove_mam_for_user_with_peer(User, Server, Peer) ->
+    LUser = jid:nodeprep(User),
+    LServer = jid:nameprep(Server),
+    try jid:decode(Peer) of
+	Jid ->
+	    Mod = gen_mod:db_mod(LUser, ?MODULE),
+	    case Mod:remove_from_archive(LUser, LServer, Jid) of
+		ok ->
+		    {ok, <<"MAM archive removed">>};
+		{error, Bin} when is_binary(Bin) ->
+		    {error, Bin};
+		{error, _} ->
+		    {error, <<"Db returned error">>}
+	    end
+    catch _:_ ->
+	{error, <<"Invalid peer JID">>}
+    end.
+
 -spec get_room_config([muc_roomconfig:property()], mod_muc_room:state(),
 		      jid(), binary()) -> [muc_roomconfig:property()].
 get_room_config(Fields, RoomState, _From, _Lang) ->
@@ -1080,7 +1112,26 @@ get_commands_spec() ->
                                      "Days to keep messages"],
 			args_example = [<<"all">>, 31],
 			args = [{type, binary}, {days, integer}],
-			result = {res, rescode}}].
+			result = {res, rescode}},
+     #ejabberd_commands{name = remove_mam_for_user, tags = [mam],
+			desc = "Remove mam archive for user",
+			module = ?MODULE, function = remove_mam_for_user,
+			args = [{user, binary}, {server, binary}],
+			args_desc = ["Username", "Server"],
+			args_example = [<<"bob">>, <<"example.com">>],
+			result = {res, restuple},
+			result_desc = "Result tuple",
+			result_example = {ok, <<"MAM archive removed">>}},
+     #ejabberd_commands{name = remove_mam_for_user_with_peer, tags = [mam],
+			desc = "Remove mam archive for user with peer",
+			module = ?MODULE, function = remove_mam_for_user_with_peer,
+			args = [{user, binary}, {server, binary}, {with, binary}],
+			args_desc = ["Username", "Server", "Peer"],
+			args_example = [<<"bob">>, <<"example.com">>, <<"anne@example.com">>],
+			result = {res, restuple},
+			result_desc = "Result tuple",
+			result_example = {ok, <<"MAM archive removed">>}}
+	].
 
 mod_opt_type(assume_mam_usage) ->
     fun (B) when is_boolean(B) -> B end;
diff --git a/src/mod_mam_mnesia.erl b/src/mod_mam_mnesia.erl
index ed6967a0..53d390cb 100644
--- a/src/mod_mam_mnesia.erl
+++ b/src/mod_mam_mnesia.erl
@@ -28,7 +28,7 @@
 
 %% API
 -export([init/2, remove_user/2, remove_room/3, delete_old_messages/3,
-	 extended_fields/0, store/8, write_prefs/4, get_prefs/2, select/6]).
+	 extended_fields/0, store/8, write_prefs/4, get_prefs/2, select/6, remove_from_archive/3]).
 
 -include_lib("stdlib/include/ms_transform.hrl").
 -include("xmpp.hrl").
@@ -67,6 +67,25 @@ remove_user(LUser, LServer) ->
 remove_room(_LServer, LName, LHost) ->
     remove_user(LName, LHost).
 
+remove_from_archive(LUser, LServer, none) ->
+    US = {LUser, LServer},
+    case mnesia:transaction(fun () -> mnesia:delete({archive_msg, US}) end) of
+	{atomic, _} -> ok;
+	{aborted, Reason} -> {error, Reason}
+    end;
+remove_from_archive(LUser, LServer, WithJid) ->
+    US = {LUser, LServer},
+    Peer = jid:remove_resource(jid:split(WithJid)),
+    F = fun () ->
+	    Msgs = mnesia:match_object(archive_msg,
+				       #archive_msg{us = US, bare_peer = Peer, _ = '_'}),
+	    lists:foreach(fun mnesia:delete_object/1, Msgs)
+	end,
+    case mnesia:transaction(F) of
+	{atomic, _} -> ok;
+	{aborted, Reason} -> {error, Reason}
+    end.
+
 delete_old_messages(global, TimeStamp, Type) ->
     mnesia:change_table_copy_type(archive_msg, node(), disc_copies),
     Result = delete_old_user_messages(mnesia:dirty_first(archive_msg), TimeStamp, Type),
diff --git a/src/mod_mam_sql.erl b/src/mod_mam_sql.erl
index 1609ff79..37ea8dc6 100644
--- a/src/mod_mam_sql.erl
+++ b/src/mod_mam_sql.erl
@@ -30,7 +30,7 @@
 
 %% API
 -export([init/2, remove_user/2, remove_room/3, delete_old_messages/3,
-	 extended_fields/0, store/8, write_prefs/4, get_prefs/2, select/6, export/1]).
+	 extended_fields/0, store/8, write_prefs/4, get_prefs/2, select/6, export/1, remove_from_archive/3]).
 
 -include_lib("stdlib/include/ms_transform.hrl").
 -include("xmpp.hrl").
@@ -56,6 +56,20 @@ remove_room(LServer, LName, LHost) ->
     LUser = jid:encode({LName, LHost, <<>>}),
     remove_user(LUser, LServer).
 
+remove_from_archive(LUser, LServer, none) ->
+    case ejabberd_sql:sql_query(LServer,
+				?SQL("delete from archive where username=%(LUser)s and %(LServer)H")) of
+	{error, Reason} -> {error, Reason};
+	_ -> ok
+    end;
+remove_from_archive(LUser, LServer, WithJid) ->
+    Peer = jid:encode(jid:remove_resource(WithJid)),
+    case ejabberd_sql:sql_query(LServer,
+				?SQL("delete from archive where username=%(LUser)s and %(LServer)H and bare_peer=%(Peer)s")) of
+	{error, Reason} -> {error, Reason};
+	_ -> ok
+    end.
+
 delete_old_messages(ServerHost, TimeStamp, Type) ->
     TS = now_to_usec(TimeStamp),
     case Type of