Improve handling of mnesia schema

1 files changed, 169 insertions, 0 deletions

diff --git a/src/ejabberd_mnesia.erl b/src/ejabberd_mnesia.erl
new file mode 100644
index 000000000..f244d518f
--- /dev/null
+++ b/src/ejabberd_mnesia.erl
@@ -0,0 +1,169 @@
+%%%----------------------------------------------------------------------
+%%% File    : mnesia_mnesia.erl
+%%% Author  : Christophe Romain <christophe.romain@process-one.net>
+%%% Purpose : Handle configurable mnesia schema
+%%% Created : 17 Nov 2016 by Christophe Romain <christophe.romain@process-one.net>
+%%%
+%%%
+%%% ejabberd, Copyright (C) 2002-2016   ProcessOne
+%%%
+%%% This program is free software; you can redistribute it and/or
+%%% modify it under the terms of the GNU General Public License as
+%%% published by the Free Software Foundation; either version 2 of the
+%%% License, or (at your option) any later version.
+%%%
+%%% This program is distributed in the hope that it will be useful,
+%%% but WITHOUT ANY WARRANTY; without even the implied warranty of
+%%% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+%%% General Public License for more details.
+%%%
+%%% You should have received a copy of the GNU General Public License along
+%%% with this program; if not, write to the Free Software Foundation, Inc.,
+%%% 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
+%%%
+%%%----------------------------------------------------------------------
+
+%%% This module should be used everywhere ejabberd creates a mnesia table
+%%% to make the schema customizable without code change
+%%% Just apply this change in ejabberd modules
+%%% s/ejabberd_mnesia:create(?MODULE, /ejabberd_mnesia:create(?MODULE, /
+
+-module(ejabberd_mnesia).
+-author('christophe.romain@process-one.net').
+-export([create/3, reset/2, update/2]).
+
+-define(STORAGE_TYPES, [disc_copies, disc_only_copies, ram_copies]).
+-define(NEED_RESET, [local_content, type]).
+
+create(Module, Name, TabDef) ->
+    Schema = schema(Module, Name, TabDef),
+    {attributes, Attrs} = lists:keyfind(attributes, 1, Schema),
+    case catch mnesia:table_info(Name, attributes) of
+	{'EXIT', _} ->
+	    ejabberd_mnesia:create(?MODULE, Name, Schema);
+	Attrs ->
+	    case need_reset(TabDef, Schema) of
+		true -> reset(Name, Schema);
+		false -> update(Name, Schema)
+	    end;
+	OldAttrs ->
+	    Fun = case lists:member({transform,1}, Module:module_info(exports)) of
+		true -> fun(Old) -> Module:transform(Old) end;
+		false -> fun(Old) -> transform(OldAttrs, Attrs, Old) end
+	    end,
+	    mnesia:transform_table(Name, Fun, Attrs)
+    end.
+
+reset(Name, TabDef) ->
+    mnesia:delete_table(Name),
+    ejabberd_mnesia:create(?MODULE, Name, TabDef).
+
+update(Name, TabDef) ->
+    Storage = mnesia:table_info(Name, storage_type),
+    NewStorage = lists:foldl(
+		   fun({Key, _}, Acc) ->
+			   case lists:member(Key, ?STORAGE_TYPES) of
+			       true -> Key;
+			       false -> Acc
+			   end
+		   end, Storage, TabDef),
+    R1 = if Storage=/=NewStorage ->
+	    mnesia:change_table_copy_type(Name, node(), NewStorage);
+       true ->
+	    {atomic, ok}
+    end,
+    Indexes = mnesia:table_info(Name, index),
+    NewIndexes = proplists:get_value(index, TabDef, []),
+    [mnesia:del_table_index(Name, Attr)
+     || Attr <- Indexes--NewIndexes],
+    R2 = [mnesia:add_table_index(Name, Attr)
+	  || Attr <- NewIndexes--Indexes],
+    lists:foldl(
+      fun({atomic, ok}, Acc) -> Acc;
+	 (Error, _Acc) -> Error
+      end, {atomic, ok}, [R1|R2]).
+
+%
+% utilities
+%
+
+schema(Module, Name, TabDef) ->
+    case parse(Module) of
+	{ok, CustomDefs} ->
+	    case lists:keyfind(Name, 1, CustomDefs) of
+		{Name, CustomDef} -> merge(TabDef, CustomDef);
+		_ -> TabDef
+	    end;
+	_ ->
+	    TabDef
+    end.
+
+merge(TabDef, CustomDef) ->
+    {CustomKeys, _} = lists:unzip(CustomDef),
+    CleanDef = lists:foldl(
+		fun(Elem, Acc) ->
+		    case lists:member(Elem, ?STORAGE_TYPES) of
+			true ->
+			    lists:foldl(
+			      fun(Key, CleanAcc) ->
+				      lists:keydelete(Key, 1, CleanAcc)
+			      end, Acc, ?STORAGE_TYPES);
+			false ->
+			    Acc
+		    end
+		end, TabDef, CustomKeys),
+    lists:ukeymerge(1,
+		   lists:ukeysort(1, CustomDef),
+		   lists:ukeysort(1, CleanDef)).
+
+parse(Module) ->
+    Path = case os:getenv("EJABBERD_SCHEMA_PATH") of
+		false ->
+		    case code:priv_dir(ejabberd) of
+			{error, _} -> "schema";  % $SPOOL_DIR/schema
+			Priv -> filename:join(Priv, "schema")
+		    end;
+		CustomDir ->
+		    CustomDir
+	    end,
+    File = filename:join(Path, atom_to_list(Module)++".mnesia"),
+    case file:consult(File) of
+	{ok, Terms} -> parse(Terms, []);
+	Error -> Error
+    end.
+
+parse([], Acc) ->
+    {ok, lists:reverse(Acc)};
+parse([{Name, Storage, TabDef}|Tail], Acc)
+  when is_atom(Name),
+       is_atom(Storage),
+       is_list(TabDef) ->
+    NewDef = case lists:member(Storage, ?STORAGE_TYPES) of
+	true -> [{Storage, [node()]} | TabDef];
+	false -> TabDef
+    end,
+    parse(Tail, [{Name, NewDef} | Acc]);
+parse([Other|_], _) ->
+    {error, {invalid, Other}}.
+
+need_reset(FromDef, ToDef) ->
+    ValuesF = [lists:keyfind(Key, 1, FromDef) || Key <- ?NEED_RESET],
+    ValuesT = [lists:keyfind(Key, 1, ToDef) || Key <- ?NEED_RESET],
+    lists:foldl(
+      fun({Val, Val}, Acc) -> Acc;
+	 ({_, false}, Acc) -> Acc;
+	 ({_, _}, _) -> true
+      end, false, lists:zip(ValuesF, ValuesT)).
+
+transform(OldAttrs, Attrs, Old) ->
+    [Name|OldValues] = tuple_to_list(Old),
+    Before = lists:zip(OldAttrs, OldValues),
+    After = lists:foldl(
+	      fun(Attr, Acc) ->
+		      case lists:keyfind(Attr, 1, Before) of
+			  false -> [{Attr, undefined}|Acc];
+			  Value -> [Value|Acc]
+		      end
+	      end, [], lists:reverse(Attrs)),
+    {Attrs, NewRecord} = lists:unzip(After),
+    list_to_tuple([Name|NewRecord]).