Improve detection of databases supported by modules (#1092)

20 files changed, 117 insertions, 81 deletions

diff --git a/src/ejabberd_auth.erl b/src/ejabberd_auth.erl
index 927abdac2..6b7f537c0 100644
--- a/src/ejabberd_auth.erl
+++ b/src/ejabberd_auth.erl
@@ -428,7 +428,7 @@ auth_modules() ->
 %% Return the list of authenticated modules for a given host
 auth_modules(Server) ->
     LServer = jid:nameprep(Server),
-    Default = gen_mod:default_db(LServer),
+    Default = ejabberd_config:default_db(LServer, ?MODULE),
     Methods = ejabberd_config:get_option(
                 {auth_method, LServer}, opt_type(auth_method), [Default]),
     [jlib:binary_to_atom(<<"ejabberd_auth_",
@@ -448,15 +448,9 @@ import(Server, riak, Passwd) ->
 import(_, _, _) ->
     pass.
 
--spec v_auth_method(atom()) -> atom().
-
-v_auth_method(odbc) -> sql;
-v_auth_method(internal) -> mnesia;
-v_auth_method(A) when is_atom(A) -> A.
-
 opt_type(auth_method) ->
     fun (V) when is_list(V) ->
-	    lists:map(fun v_auth_method/1, V);
-	(V) -> [v_auth_method(V)]
+	    lists:map(fun(M) -> ejabberd_config:v_db(?MODULE, M) end, V);
+	(V) -> [ejabberd_config:v_db(?MODULE, V)]
     end;
 opt_type(_) -> [auth_method].
diff --git a/src/ejabberd_config.erl b/src/ejabberd_config.erl
index 7fccbc744..87f465e13 100644
--- a/src/ejabberd_config.erl
+++ b/src/ejabberd_config.erl
@@ -34,8 +34,8 @@
          get_vh_by_auth_method/1, is_file_readable/1,
          get_version/0, get_myhosts/0, get_mylang/0,
          prepare_opt_val/4, convert_table_to_binary/5,
-         transform_options/1, collect_options/1,
-         convert_to_yaml/1, convert_to_yaml/2,
+         transform_options/1, collect_options/1, default_db/2,
+         convert_to_yaml/1, convert_to_yaml/2, v_db/2,
          env_binary_to_list/2, opt_type/1, may_hide_data/1]).
 
 -export([start/2]).
@@ -674,7 +674,16 @@ rename_option(Option) ->
 
 change_val(auth_method, Val) ->
     prepare_opt_val(auth_method, Val,
-		    ejabberd_auth:opt_type(auth_method), [mnesia]);
+		    fun(V) ->
+			    L = if is_list(V) -> V;
+				   true -> [V]
+				end,
+			    lists:map(
+			      fun(odbc) -> sql;
+				 (internal) -> mnesia;
+				 (A) when is_atom(A) -> A
+			      end, L)
+		    end, [mnesia]);
 change_val(_Opt, Val) ->
     Val.
 
@@ -806,9 +815,57 @@ get_option(Opt, F, Default) ->
             end
     end.
 
+init_module_db_table(Modules) ->
+    catch ets:new(module_db, [named_table, public, bag]),
+    %% Dirty hack for mod_pubsub
+    ets:insert(module_db, {mod_pubsub, mnesia}),
+    ets:insert(module_db, {mod_pubsub, sql}),
+    lists:foreach(
+      fun(M) ->
+	      case re:split(atom_to_list(M), "_", [{return, list}]) of
+		  [_] ->
+		      ok;
+		  Parts ->
+		      [Suffix|T] = lists:reverse(Parts),
+		      BareMod = string:join(lists:reverse(T), "_"),
+		      ets:insert(module_db, {list_to_atom(BareMod),
+					     list_to_atom(Suffix)})
+	      end
+      end, Modules).
+
+-spec v_db(module(), atom()) -> atom().
+
+v_db(Mod, internal) -> v_db(Mod, mnesia);
+v_db(Mod, odbc) -> v_db(Mod, sql);
+v_db(Mod, Type) ->
+    case ets:match_object(module_db, {Mod, Type}) of
+	[_|_] -> Type;
+	[] -> erlang:error(badarg)
+    end.
+
+-spec default_db(binary(), module()) -> atom().
+
+default_db(Host, Module) ->
+    case ejabberd_config:get_option(
+	   {default_db, Host}, fun(T) when is_atom(T) -> T end) of
+	undefined ->
+	    mnesia;
+	DBType ->
+	    try
+		v_db(Module, DBType)
+	    catch error:badarg ->
+		    ?WARNING_MSG("Module '~s' doesn't support database '~s' "
+				 "defined in option 'default_db', using "
+				 "'mnesia' as fallback", [Module, DBType]),
+		    mnesia
+	    end
+    end.
+
 get_modules_with_options() ->
     {ok, Mods} = application:get_key(ejabberd, modules),
     ExtMods = [Name || {Name, _Details} <- ext_mod:installed()],
+    AllMods = [?MODULE|ExtMods++Mods],
+    init_module_db_table(AllMods),
     lists:foldl(
       fun(Mod, D) ->
 	      case catch Mod:opt_type('') of
@@ -820,7 +877,7 @@ get_modules_with_options() ->
 		  {'EXIT', {undef, _}} ->
 		      D
 	      end
-      end, dict:new(), [?MODULE|ExtMods++Mods]).
+      end, dict:new(), AllMods).
 
 validate_opts(#state{opts = Opts} = State) ->
     ModOpts = get_modules_with_options(),
diff --git a/src/ejabberd_riak_sup.erl b/src/ejabberd_riak_sup.erl
index af811441b..ad65ecf80 100644
--- a/src/ejabberd_riak_sup.erl
+++ b/src/ejabberd_riak_sup.erl
@@ -70,8 +70,8 @@ is_riak_configured(Host) ->
 		{modules, Host},
 		fun(L) when is_list(L) -> L end, []),
     ModuleWithRiakDBConfigured = lists:any(
-				   fun({_Module, Opts}) ->
-					   gen_mod:db_type(Host, Opts) == riak
+				   fun({Module, Opts}) ->
+					   gen_mod:db_type(Host, Opts, Module) == riak
 				   end, Modules),
     ServerConfigured or PortConfigured
 	or AuthConfigured or ModuleWithRiakDBConfigured.
diff --git a/src/ejabberd_sm.erl b/src/ejabberd_sm.erl
index 9a93e6d97..20b0658d5 100644
--- a/src/ejabberd_sm.erl
+++ b/src/ejabberd_sm.erl
@@ -733,13 +733,10 @@ force_update_presence({LUser, LServer}) ->
 -spec get_sm_backend(binary()) -> module().
 
 get_sm_backend(Host) ->
-    DBType = ejabberd_config:get_option({sm_db_type, Host},
-					fun(mnesia) -> mnesia;
-					   (internal) -> mnesia;
-					   (odbc) -> sql;
-					   (sql) -> sql;
-					   (redis) -> redis
-					end, mnesia),
+    DBType = ejabberd_config:get_option(
+	       {sm_db_type, Host},
+	       fun(T) -> ejabberd_config:v_db(?MODULE, T) end,
+	       mnesia),
     list_to_atom("ejabberd_sm_" ++ atom_to_list(DBType)).
 
 -spec get_sm_backends() -> [module()].
@@ -812,11 +809,5 @@ kick_user(User, Server) ->
 make_sid() ->
     {p1_time_compat:unique_timestamp(), self()}.
 
-opt_type(sm_db_type) ->
-    fun (mnesia) -> mnesia;
-	(internal) -> mnesia;
-	(sql) -> sql;
-	(odbc) -> sql;
-	(redis) -> redis
-    end;
+opt_type(sm_db_type) -> fun(T) -> ejabberd_config:v_db(?MODULE, T) end;
 opt_type(_) -> [sm_db_type].
diff --git a/src/gen_mod.erl b/src/gen_mod.erl
index 26e662dc6..efbfc087f 100644
--- a/src/gen_mod.erl
+++ b/src/gen_mod.erl
@@ -31,12 +31,12 @@
 
 -export([start/0, start_module/2, start_module/3,
 	 stop_module/2, stop_module_keep_config/2, get_opt/3,
-	 get_opt/4, get_opt_host/3, db_type/1, db_type/2,
+	 get_opt/4, get_opt_host/3, db_type/2, db_type/3,
 	 get_module_opt/4, get_module_opt/5, get_module_opt_host/3,
 	 loaded_modules/1, loaded_modules_with_opts/1,
 	 get_hosts/2, get_module_proc/2, is_loaded/2,
 	 start_modules/0, start_modules/1, stop_modules/0, stop_modules/1,
-	 default_db/1, v_db/1, opt_type/1, db_mod/2, db_mod/3]).
+	 opt_type/1, db_mod/2, db_mod/3]).
 
 %%-export([behaviour_info/1]).
 
@@ -295,44 +295,46 @@ validate_opts(Module, Opts) ->
 	      false
       end, Opts).
 
--spec v_db(db_type() | internal) -> db_type().
-
-v_db(odbc) -> sql;
-v_db(sql) -> sql;
-v_db(internal) -> mnesia;
-v_db(mnesia) -> mnesia;
-v_db(riak) -> riak.
-
--spec db_type(opts()) -> db_type().
-
-db_type(Opts) ->
-    db_type(global, Opts).
-
--spec db_type(binary() | global, atom() | opts()) -> db_type().
+-spec db_type(binary() | global, module()) -> db_type();
+	     (opts(), module()) -> db_type().
 
+db_type(Opts, Module) when is_list(Opts) ->
+    db_type(global, Opts, Module);
 db_type(Host, Module) when is_atom(Module) ->
-    get_module_opt(Host, Module, db_type, fun v_db/1, default_db(Host));
-db_type(Host, Opts) when is_list(Opts) ->
-    get_opt(db_type, Opts, fun v_db/1, default_db(Host)).
-
--spec default_db(binary() | global) -> db_type().
+    case Module:mod_opt_type(db_type) of
+	F when is_function(F) ->
+	    case get_module_opt(Host, Module, db_type, F) of
+		undefined -> ejabberd_config:default_db(Host, Module);
+		Type -> Type
+	    end;
+	_ ->
+	    undefined
+    end.
 
-default_db(Host) ->
-    ejabberd_config:get_option({default_db, Host}, fun v_db/1, mnesia).
+-spec db_type(binary(), opts(), module()) -> db_type().
+
+db_type(Host, Opts, Module) ->
+    case Module:mod_opt_type(db_type) of
+	F when is_function(F) ->
+	    case get_opt(db_type, Opts, F) of
+		undefined -> ejabberd_config:default_db(Host, Module);
+		Type -> Type
+	    end;
+	_ ->
+	    undefined
+    end.
 
 -spec db_mod(binary() | global | db_type(), module()) -> module().
 
-db_mod(odbc, Module) -> list_to_atom(atom_to_list(Module) ++ "_sql");
-db_mod(sql, Module) -> list_to_atom(atom_to_list(Module) ++ "_sql");
-db_mod(mnesia, Module) -> list_to_atom(atom_to_list(Module) ++ "_mnesia");
-db_mod(riak, Module) -> list_to_atom(atom_to_list(Module) ++ "_riak");
+db_mod(Type, Module) when is_atom(Type) ->
+    list_to_atom(atom_to_list(Module) ++ "_" ++ atom_to_list(Type));
 db_mod(Host, Module) when is_binary(Host) orelse Host == global ->
     db_mod(db_type(Host, Module), Module).
 
 -spec db_mod(binary() | global, opts(), module()) -> module().
 
 db_mod(Host, Opts, Module) when is_list(Opts) ->
-    db_mod(db_type(Host, Opts), Module).
+    db_mod(db_type(Host, Opts, Module), Module).
 
 -spec loaded_modules(binary()) -> [atom()].
 
@@ -380,6 +382,6 @@ get_module_proc(Host, Base) ->
 is_loaded(Host, Module) ->
     ets:member(ejabberd_modules, {Module, Host}).
 
-opt_type(default_db) -> fun v_db/1;
+opt_type(default_db) -> fun(T) when is_atom(T) -> T end;
 opt_type(modules) -> fun (L) when is_list(L) -> L end;
 opt_type(_) -> [default_db, modules].
diff --git a/src/mod_announce.erl b/src/mod_announce.erl
index d7251c50b..9a9e665f5 100644
--- a/src/mod_announce.erl
+++ b/src/mod_announce.erl
@@ -921,5 +921,5 @@ import(LServer, DBType, LA) ->
 
 mod_opt_type(access) ->
     fun (A) when is_atom(A) -> A end;
-mod_opt_type(db_type) -> fun gen_mod:v_db/1;
+mod_opt_type(db_type) -> fun(T) -> ejabberd_config:v_db(?MODULE, T) end;
 mod_opt_type(_) -> [access, db_type].
diff --git a/src/mod_caps.erl b/src/mod_caps.erl
index 3d5c360a8..966a9baa6 100644
--- a/src/mod_caps.erl
+++ b/src/mod_caps.erl
@@ -651,6 +651,6 @@ mod_opt_type(cache_life_time) ->
     fun (I) when is_integer(I), I > 0 -> I end;
 mod_opt_type(cache_size) ->
     fun (I) when is_integer(I), I > 0 -> I end;
-mod_opt_type(db_type) -> fun gen_mod:v_db/1;
+mod_opt_type(db_type) -> fun(T) -> ejabberd_config:v_db(?MODULE, T) end;
 mod_opt_type(_) ->
     [cache_life_time, cache_size, db_type].
diff --git a/src/mod_carboncopy.erl b/src/mod_carboncopy.erl
index ebf1d0b0f..bb20bd2f9 100644
--- a/src/mod_carboncopy.erl
+++ b/src/mod_carboncopy.erl
@@ -279,8 +279,5 @@ list(User, Server) ->
     Mod:list(User, Server).
 
 mod_opt_type(iqdisc) -> fun gen_iq_handler:check_type/1;
-mod_opt_type(db_type) ->
-    fun(internal) -> mnesia;
-       (mnesia) -> mnesia
-    end;
+mod_opt_type(db_type) -> fun(T) -> ejabberd_config:v_db(?MODULE, T) end;
 mod_opt_type(_) -> [db_type, iqdisc].
diff --git a/src/mod_irc.erl b/src/mod_irc.erl
index e0c658dea..f6487a1a9 100644
--- a/src/mod_irc.erl
+++ b/src/mod_irc.erl
@@ -1253,7 +1253,7 @@ import(LServer, DBType, Data) ->
 
 mod_opt_type(access) ->
     fun (A) when is_atom(A) -> A end;
-mod_opt_type(db_type) -> fun gen_mod:v_db/1;
+mod_opt_type(db_type) -> fun(T) -> ejabberd_config:v_db(?MODULE, T) end;
 mod_opt_type(default_encoding) ->
     fun iolist_to_binary/1;
 mod_opt_type(host) -> fun iolist_to_binary/1;
diff --git a/src/mod_last.erl b/src/mod_last.erl
index 1af1847b3..c39681f69 100644
--- a/src/mod_last.erl
+++ b/src/mod_last.erl
@@ -244,7 +244,7 @@ transform_options({node_start, {_, _, _} = Now}, Opts) ->
 transform_options(Opt, Opts) ->
     [Opt|Opts].
 
-mod_opt_type(db_type) -> fun gen_mod:v_db/1;
+mod_opt_type(db_type) -> fun(T) -> ejabberd_config:v_db(?MODULE, T) end;
 mod_opt_type(iqdisc) -> fun gen_iq_handler:check_type/1;
 mod_opt_type(_) -> [db_type, iqdisc].
 
diff --git a/src/mod_mam.erl b/src/mod_mam.erl
index f84519a08..d5ee9bf86 100644
--- a/src/mod_mam.erl
+++ b/src/mod_mam.erl
@@ -1017,12 +1017,7 @@ mod_opt_type(cache_life_time) ->
     fun (I) when is_integer(I), I > 0 -> I end;
 mod_opt_type(cache_size) ->
     fun (I) when is_integer(I), I > 0 -> I end;
-mod_opt_type(db_type) ->
-    fun(sql) -> sql;
-       (odbc) -> sql;
-       (internal) -> mnesia;
-       (mnesia) -> mnesia
-    end;
+mod_opt_type(db_type) -> fun(T) -> ejabberd_config:v_db(?MODULE, T) end;
 mod_opt_type(default) ->
     fun (always) -> always;
 	(never) -> never;
diff --git a/src/mod_muc.erl b/src/mod_muc.erl
index 6aa186318..3d098e0d5 100644
--- a/src/mod_muc.erl
+++ b/src/mod_muc.erl
@@ -932,7 +932,7 @@ mod_opt_type(access_create) ->
     fun (A) when is_atom(A) -> A end;
 mod_opt_type(access_persistent) ->
     fun (A) when is_atom(A) -> A end;
-mod_opt_type(db_type) -> fun gen_mod:v_db/1;
+mod_opt_type(db_type) -> fun(T) -> ejabberd_config:v_db(?MODULE, T) end;
 mod_opt_type(default_room_options) ->
     fun (L) when is_list(L) -> L end;
 mod_opt_type(history_size) ->
diff --git a/src/mod_offline.erl b/src/mod_offline.erl
index 356d89a67..4d8aba762 100644
--- a/src/mod_offline.erl
+++ b/src/mod_offline.erl
@@ -867,7 +867,7 @@ import(LServer, DBType, Data) ->
 
 mod_opt_type(access_max_user_messages) ->
     fun (A) -> A end;
-mod_opt_type(db_type) -> fun gen_mod:v_db/1;
+mod_opt_type(db_type) -> fun(T) -> ejabberd_config:v_db(?MODULE, T) end;
 mod_opt_type(store_empty_body) ->
     fun (V) when is_boolean(V) -> V;
         (unless_chat_state) -> unless_chat_state
diff --git a/src/mod_privacy.erl b/src/mod_privacy.erl
index 413dcb52a..ad13c27cd 100644
--- a/src/mod_privacy.erl
+++ b/src/mod_privacy.erl
@@ -593,6 +593,6 @@ import(LServer, DBType, Data) ->
     Mod = gen_mod:db_mod(DBType, ?MODULE),
     Mod:import(LServer, Data).
 
-mod_opt_type(db_type) -> fun gen_mod:v_db/1;
+mod_opt_type(db_type) -> fun(T) -> ejabberd_config:v_db(?MODULE, T) end;
 mod_opt_type(iqdisc) -> fun gen_iq_handler:check_type/1;
 mod_opt_type(_) -> [db_type, iqdisc].
diff --git a/src/mod_private.erl b/src/mod_private.erl
index 029789e63..38e42ca4c 100644
--- a/src/mod_private.erl
+++ b/src/mod_private.erl
@@ -173,6 +173,6 @@ import(LServer, DBType, PD) ->
     Mod = gen_mod:db_mod(DBType, ?MODULE),
     Mod:import(LServer, PD).
 
-mod_opt_type(db_type) -> fun gen_mod:v_db/1;
+mod_opt_type(db_type) -> fun(T) -> ejabberd_config:v_db(?MODULE, T) end;
 mod_opt_type(iqdisc) -> fun gen_iq_handler:check_type/1;
 mod_opt_type(_) -> [db_type, iqdisc].
diff --git a/src/mod_pubsub.erl b/src/mod_pubsub.erl
index a86155af8..a801fa984 100644
--- a/src/mod_pubsub.erl
+++ b/src/mod_pubsub.erl
@@ -4470,7 +4470,7 @@ purge_offline(Host, LJID, Node) ->
 
 mod_opt_type(access_createnode) ->
     fun (A) when is_atom(A) -> A end;
-mod_opt_type(db_type) -> fun gen_mod:v_db/1;
+mod_opt_type(db_type) -> fun(T) -> ejabberd_config:v_db(?MODULE, T) end;
 mod_opt_type(host) -> fun iolist_to_binary/1;
 mod_opt_type(ignore_pep_from_offline) ->
     fun (A) when is_boolean(A) -> A end;
diff --git a/src/mod_roster.erl b/src/mod_roster.erl
index 16354dd8f..b3a627f7c 100644
--- a/src/mod_roster.erl
+++ b/src/mod_roster.erl
@@ -1236,7 +1236,7 @@ import(LServer, DBType, R) ->
 
 mod_opt_type(access) ->
     fun (A) when is_atom(A) -> A end;
-mod_opt_type(db_type) -> fun gen_mod:v_db/1;
+mod_opt_type(db_type) -> fun(T) -> ejabberd_config:v_db(?MODULE, T) end;
 mod_opt_type(iqdisc) -> fun gen_iq_handler:check_type/1;
 mod_opt_type(managers) ->
     fun (B) when is_list(B) -> B end;
diff --git a/src/mod_shared_roster.erl b/src/mod_shared_roster.erl
index 6670cf77b..76a619c9b 100644
--- a/src/mod_shared_roster.erl
+++ b/src/mod_shared_roster.erl
@@ -1120,5 +1120,5 @@ import(LServer, DBType, Data) ->
     Mod = gen_mod:db_mod(DBType, ?MODULE),
     Mod:import(LServer, Data).
 
-mod_opt_type(db_type) -> fun gen_mod:v_db/1;
+mod_opt_type(db_type) -> fun(T) -> ejabberd_config:v_db(?MODULE, T) end;
 mod_opt_type(_) -> [db_type].
diff --git a/src/mod_vcard.erl b/src/mod_vcard.erl
index e5f5d9e3c..5e042528b 100644
--- a/src/mod_vcard.erl
+++ b/src/mod_vcard.erl
@@ -596,7 +596,7 @@ import(LServer, DBType, VCard) ->
 
 mod_opt_type(allow_return_all) ->
     fun (B) when is_boolean(B) -> B end;
-mod_opt_type(db_type) -> fun gen_mod:v_db/1;
+mod_opt_type(db_type) -> fun(T) -> ejabberd_config:v_db(?MODULE, T) end;
 mod_opt_type(host) -> fun iolist_to_binary/1;
 mod_opt_type(iqdisc) -> fun gen_iq_handler:check_type/1;
 mod_opt_type(matches) ->
diff --git a/src/mod_vcard_xupdate.erl b/src/mod_vcard_xupdate.erl
index 198312c36..041b0b64c 100644
--- a/src/mod_vcard_xupdate.erl
+++ b/src/mod_vcard_xupdate.erl
@@ -133,5 +133,5 @@ import(LServer, DBType, LA) ->
     Mod = gen_mod:db_mod(DBType, ?MODULE),
     Mod:import(LServer, LA).
 
-mod_opt_type(db_type) -> fun gen_mod:v_db/1;
+mod_opt_type(db_type) -> fun(T) -> ejabberd_config:v_db(?MODULE, T) end;
 mod_opt_type(_) -> [db_type].