From 3029e84faf5bcdf1c01cc0376b22b2da541c0cc8 Mon Sep 17 00:00:00 2001
From: Christophe Romain <christophe.romain@process-one.net>
Date: Wed, 18 Jan 2017 16:53:36 +0100
Subject: [PATCH] Improve check of mnesia schema opearations

---
 src/ejabberd_mnesia.erl | 91 ++++++++++++++++++++++++-----------------
 1 file changed, 54 insertions(+), 37 deletions(-)

diff --git a/src/ejabberd_mnesia.erl b/src/ejabberd_mnesia.erl
index 0a234a9d0..e38168594 100644
--- a/src/ejabberd_mnesia.erl
+++ b/src/ejabberd_mnesia.erl
@@ -35,12 +35,16 @@
 -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),
+-include("logger.hrl").
+
+create(Module, Name, TabDef)
+  when is_atom(Module), is_atom(Name), is_list(TabDef) ->
+    Path = os:getenv("EJABBERD_SCHEMA_PATH"),
+    Schema = schema(Path, Module, Name, TabDef),
     {attributes, Attrs} = lists:keyfind(attributes, 1, Schema),
     case catch mnesia:table_info(Name, attributes) of
 	{'EXIT', _} ->
-	    mnesia:create_table(Name, Schema);
+	    mnesia_op(create, [Name, TabDef]);
 	Attrs ->
 	    case need_reset(TabDef, Schema) of
 		true -> reset(Name, Schema);
@@ -51,18 +55,23 @@ create(Module, Name, TabDef) ->
 		true -> fun(Old) -> Module:transform(Old) end;
 		false -> fun(Old) -> transform(OldAttrs, Attrs, Old) end
 	    end,
-	    mnesia:transform_table(Name, Fun, Attrs)
+	    mnesia_op(transform_table, [Name, Fun, Attrs])
     end.
 
-reset(Name, TabDef) ->
-    mnesia:delete_table(Name),
-    ejabberd_mnesia:create(?MODULE, Name, TabDef).
+reset(Name, TabDef)
+  when is_atom(Name), is_list(TabDef) ->
+    mnesia_op(delete_table, [Name]),
+    mnesia_op(create, [Name, TabDef]).
 
-update(Name, TabDef) ->
+update(Name, TabDef)
+  when is_atom(Name), is_list(TabDef) ->
     {attributes, Attrs} = lists:keyfind(attributes, 1, TabDef),
     update(Name, Attrs, TabDef).
 update(Name, Attrs, TabDef) ->
-    Storage = mnesia:table_info(Name, storage_type),
+    Storage = case catch mnesia:table_info(Name, storage_type) of
+	{'EXIT', _} -> unknown;
+	Type -> Type
+	end,
     NewStorage = lists:foldl(
 		   fun({Key, _}, Acc) ->
 			   case lists:member(Key, ?STORAGE_TYPES) of
@@ -70,34 +79,41 @@ update(Name, Attrs, TabDef) ->
 			       false -> Acc
 			   end
 		   end, Storage, TabDef),
-    R1 = if Storage=/=NewStorage ->
-	    mnesia:change_table_copy_type(Name, node(), NewStorage);
-       true ->
-	    {atomic, ok}
-    end,
+    R1 = [mnesia_op(change_table_copy_type, [Name, node(), NewStorage])
+	  || Storage=/=NewStorage],
     CurIndexes = [lists:nth(N-1, Attrs) || N<-mnesia:table_info(Name, index)],
     NewIndexes = proplists:get_value(index, TabDef, []),
-    [mnesia:del_table_index(Name, Attr)
-     || Attr <- CurIndexes--NewIndexes],
-    R2 = [mnesia:add_table_index(Name, Attr)
+    R2 = [mnesia_op(del_table_index, [Name, Attr])
+	  || Attr <- CurIndexes--NewIndexes],
+    R3 = [mnesia_op(add_table_index, [Name, Attr])
 	  || Attr <- NewIndexes--CurIndexes],
     lists:foldl(
       fun({atomic, ok}, Acc) -> Acc;
 	 (Error, _Acc) -> Error
-      end, {atomic, ok}, [R1|R2]).
+      end, {atomic, ok}, R1++R2++R3).
 
 %
 % utilities
 %
 
-schema(Module, Name, TabDef) ->
-    case parse(Module) of
+schema(false, Module, _Name, TabDef) ->
+    ?DEBUG("No custom ~s schema path", [Module]),
+    TabDef;
+schema(Path, Module, Name, TabDef) ->
+    File = filename:join(Path, atom_to_list(Module)++".mnesia"),
+    case parse(File) of
 	{ok, CustomDefs} ->
 	    case lists:keyfind(Name, 1, CustomDefs) of
-		{Name, CustomDef} -> merge(TabDef, CustomDef);
-		_ -> TabDef
+		{Name, CustomDef} ->
+		    ?INFO_MSG("Using custom ~s schema for table ~s",
+			      [Module, Name]),
+		    merge(TabDef, CustomDef);
+		_ ->
+		    TabDef
 	    end;
-	_ ->
+	{error, Error} ->
+	    ?ERROR_MSG("Can not use custom ~s schema for table ~s: ~p",
+		       [Module, Name, Error]),
 	    TabDef
     end.
 
@@ -119,24 +135,15 @@ merge(TabDef, CustomDef) ->
 		   lists:ukeysort(1, CustomDef),
 		   lists:ukeysort(1, CleanDef)).
 
-parse(Module) ->
-    case os:getenv("EJABBERD_SCHEMA_PATH") of
-	false ->
-	    {error, undefined};
-	Path ->
-	    File = filename:join(Path, atom_to_list(Module)++".mnesia"),
-	    case file:consult(File) of
-		{ok, Terms} -> parse(Terms, []);
-		Error -> Error
-	    end
+parse(File) ->
+    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) ->
+  when is_atom(Name), is_atom(Storage), is_list(TabDef) ->
     NewDef = case lists:member(Storage, ?STORAGE_TYPES) of
 	true -> [{Storage, [node()]} | TabDef];
 	false -> TabDef
@@ -166,3 +173,13 @@ transform(OldAttrs, Attrs, Old) ->
 	      end, [], lists:reverse(Attrs)),
     {Attrs, NewRecord} = lists:unzip(After),
     list_to_tuple([Name|NewRecord]).
+
+mnesia_op(Fun, Args) ->
+    case apply(mnesia, Fun, Args) of
+	{atomic, ok} ->
+	    {atomic, ok};
+	Other ->
+	    ?ERROR_MSG("failure on mnesia ~s ~p: ~p",
+		      [Fun, Args, Other]),
+	    Other
+    end.