xmpp.chapril.org-ejabberd/src/ejabberd_mnesia.erl

%%%----------------------------------------------------------------------
%%% 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-2017   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').

-behaviour(gen_server).

-export([start/0, create/3, reset/2, update/2]).
%% gen_server callbacks
-export([init/1, handle_call/3, handle_cast/2, handle_info/2,
	 terminate/2, code_change/3]).

-define(STORAGE_TYPES, [disc_copies, disc_only_copies, ram_copies]).
-define(NEED_RESET, [local_content, type]).

-include("logger.hrl").

-record(state, {tables = #{} :: map()}).

start() ->
    gen_server:start_link({local, ?MODULE}, ?MODULE, [], []).

create(Module, Name, TabDef) ->
    gen_server:call(?MODULE, {create, Module, Name, TabDef},
		    timer:seconds(60)).

init([]) ->
    ejabberd_config:env_binary_to_list(mnesia, dir),
    MyNode = node(),
    DbNodes = mnesia:system_info(db_nodes),
    case lists:member(MyNode, DbNodes) of
	true ->
	    case mnesia:system_info(extra_db_nodes) of
		[] -> mnesia:create_schema([node()]);
		_ -> ok
	    end,
	    ejabberd:start_app(mnesia, permanent),
	    ?DEBUG("Waiting for Mnesia tables synchronization...", []),
	    mnesia:wait_for_tables(mnesia:system_info(local_tables), infinity),
	    {ok, #state{}};
	false ->
	    ?CRITICAL_MSG("Node name mismatch: I'm [~s], "
			  "the database is owned by ~p", [MyNode, DbNodes]),
	    ?CRITICAL_MSG("Either set ERLANG_NODE in ejabberdctl.cfg "
			  "or change node name in Mnesia", []),
	    {stop, node_name_mismatch}
    end.

handle_call({create, Module, Name, TabDef}, _From, State) ->
    case maps:get(Name, State#state.tables, undefined) of
	{TabDef, Result} ->
	    {reply, Result, State};
	_ ->
	    Result = do_create(Module, Name, TabDef),
	    Tables = maps:put(Name, {TabDef, Result}, State#state.tables),
	    {reply, Result, #state{tables = Tables}}
    end;
handle_call(_Request, _From, State) ->
    {noreply, State}.

handle_cast(_Msg, State) ->
    {noreply, State}.

handle_info(_Info, State) ->
    {noreply, State}.

terminate(_Reason, _State) ->
    ok.

code_change(_OldVsn, State, _Extra) ->
    {ok, State}.

do_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', _} ->
	    create(Name, TabDef);
	Attrs ->
	    case need_reset(Name, Schema) of
		true -> reset(Name, Schema);
		false -> update(Name, Attrs, 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_op(transform_table, [Name, Fun, Attrs])
    end.

reset(Name, TabDef)
  when is_atom(Name), is_list(TabDef) ->
    mnesia_op(delete_table, [Name]),
    create(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 = 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
			       true -> Key;
			       false -> Acc
			   end
		   end, Storage, TabDef),
    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, []),
    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++R3).

%
% utilities
%

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} ->
		    ?INFO_MSG("Using custom ~s schema for table ~s",
			      [Module, Name]),
		    merge(TabDef, CustomDef);
		_ ->
		    TabDef
	    end;
	{error, enoent} ->
	    ?DEBUG("No custom ~s schema path", [Module]),
	    TabDef;
	{error, Error} ->
	    ?ERROR_MSG("Can not use custom ~s schema for table ~s: ~p",
		       [Module, Name, Error]),
	    TabDef
    end.

create(Name, TabDef) ->
    case mnesia_op(create_table, [Name, TabDef]) of
	{atomic, ok} ->
	    add_table_copy(Name);
	Err ->
	    Err
    end.

%% The table MUST exist, otherwise the function would fail
add_table_copy(Name) ->
    Type = mnesia:table_info(Name, storage_type),
    Nodes = mnesia:table_info(Name, Type),
    case lists:member(node(), Nodes) of
	true ->
	    {atomic, ok};
	false ->
	    mnesia_op(add_table_copy, [Name, node(), Type])
    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(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) ->
    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(Table, TabDef) ->
    ValuesF = [mnesia:table_info(Table, Key) || Key <- ?NEED_RESET],
    ValuesT = [proplists:get_value(Key, TabDef) || Key <- ?NEED_RESET],
    lists:foldl(
      fun({Val, Val}, Acc) -> Acc;
	 ({_, undefined}, 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]).

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.