1 files changed, 22 insertions, 16 deletions

diff --git a/src/mod_privacy.erl b/src/mod_privacy.erl
index f93f26f7f..9c9ec919f 100644
--- a/src/mod_privacy.erl
+++ b/src/mod_privacy.erl
@@ -43,7 +43,7 @@
 	 sql_get_privacy_list_data_by_id_t/1,
 	 sql_get_privacy_list_id_t/2,
 	 sql_set_default_privacy_list/2,
-	 sql_set_privacy_list/2]).
+	 sql_set_privacy_list/2, privacy_schema/0]).
 
 -include("ejabberd.hrl").
 -include("logger.hrl").
@@ -52,6 +52,9 @@
 
 -include("mod_privacy.hrl").
 
+privacy_schema() ->
+    {record_info(fields, privacy), #privacy{}}.
+
 start(Host, Opts) ->
     IQDisc = gen_mod:get_opt(iqdisc, Opts, fun gen_iq_handler:check_type/1,
                              one_queue),
@@ -161,7 +164,7 @@ process_lists_get(LUser, LServer, _Active, mnesia) ->
 	  {Default, LItems}
     end;
 process_lists_get(LUser, LServer, _Active, riak) ->
-    case ejabberd_riak:get(privacy, {LUser, LServer}) of
+    case ejabberd_riak:get(privacy, privacy_schema(), {LUser, LServer}) of
         {ok, #privacy{default = Default, lists = Lists}} ->
             LItems = lists:map(fun ({N, _}) ->
                                        #xmlel{name = <<"list">>,
@@ -225,7 +228,7 @@ process_list_get(LUser, LServer, Name, mnesia) ->
 	  end
     end;
 process_list_get(LUser, LServer, Name, riak) ->
-    case ejabberd_riak:get(privacy, {LUser, LServer}) of
+    case ejabberd_riak:get(privacy, privacy_schema(), {LUser, LServer}) of
         {ok, #privacy{lists = Lists}} ->
             case lists:keysearch(Name, 1, Lists) of
                 {value, {_, List}} -> List;
@@ -383,12 +386,13 @@ process_default_set(LUser, LServer, {value, Name},
     mnesia:transaction(F);
 process_default_set(LUser, LServer, {value, Name}, riak) ->
     {atomic,
-     case ejabberd_riak:get(privacy, {LUser, LServer}) of
+     case ejabberd_riak:get(privacy, privacy_schema(), {LUser, LServer}) of
          {ok, #privacy{lists = Lists} = P} ->
              case lists:keymember(Name, 1, Lists) of
                  true ->
                      ejabberd_riak:put(P#privacy{default = Name,
-                                                 lists = Lists});
+                                                 lists = Lists},
+				       privacy_schema());
                  false ->
                      not_found
              end;
@@ -418,9 +422,9 @@ process_default_set(LUser, LServer, false, mnesia) ->
     mnesia:transaction(F);
 process_default_set(LUser, LServer, false, riak) ->
     {atomic,
-     case ejabberd_riak:get(privacy, {LUser, LServer}) of
+     case ejabberd_riak:get(privacy, privacy_schema(), {LUser, LServer}) of
          {ok, R} ->
-             ejabberd_riak:put(R#privacy{default = none});
+             ejabberd_riak:put(R#privacy{default = none}, privacy_schema());
          {error, _} ->
              ok
      end};
@@ -457,7 +461,7 @@ process_active_set(LUser, LServer, Name, mnesia) ->
 	  end
     end;
 process_active_set(LUser, LServer, Name, riak) ->
-    case ejabberd_riak:get(privacy, {LUser, LServer}) of
+    case ejabberd_riak:get(privacy, privacy_schema(), {LUser, LServer}) of
         {ok, #privacy{lists = Lists}} ->
             case lists:keysearch(Name, 1, Lists) of
                 {value, {_, List}} -> List;
@@ -499,13 +503,14 @@ remove_privacy_list(LUser, LServer, Name, mnesia) ->
     mnesia:transaction(F);
 remove_privacy_list(LUser, LServer, Name, riak) ->
     {atomic,
-     case ejabberd_riak:get(privacy, {LUser, LServer}) of
+     case ejabberd_riak:get(privacy, privacy_schema(), {LUser, LServer}) of
          {ok, #privacy{default = Default, lists = Lists} = P} ->
              if Name == Default ->
                      conflict;
                 true ->
                      NewLists = lists:keydelete(Name, 1, Lists),
-                     ejabberd_riak:put(P#privacy{lists = NewLists})
+                     ejabberd_riak:put(P#privacy{lists = NewLists},
+				       privacy_schema())
              end;
          {error, _} ->
              ok
@@ -539,15 +544,16 @@ set_privacy_list(LUser, LServer, Name, List, mnesia) ->
     mnesia:transaction(F);
 set_privacy_list(LUser, LServer, Name, List, riak) ->
     {atomic,
-     case ejabberd_riak:get(privacy, {LUser, LServer}) of
+     case ejabberd_riak:get(privacy, privacy_schema(), {LUser, LServer}) of
          {ok, #privacy{lists = Lists} = P} ->
              NewLists1 = lists:keydelete(Name, 1, Lists),
              NewLists = [{Name, List} | NewLists1],
-             ejabberd_riak:put(P#privacy{lists = NewLists});
+             ejabberd_riak:put(P#privacy{lists = NewLists}, privacy_schema());
          {error, _} ->
              NewLists = [{Name, List}],
              ejabberd_riak:put(#privacy{us = {LUser, LServer},
-                                        lists = NewLists})
+                                        lists = NewLists},
+			       privacy_schema())
      end};
 set_privacy_list(LUser, LServer, Name, List, odbc) ->
     RItems = lists:map(fun item_to_raw/1, List),
@@ -734,7 +740,7 @@ get_user_list(_, LUser, LServer, mnesia) ->
       _ -> {none, []}
     end;
 get_user_list(_, LUser, LServer, riak) ->
-    case ejabberd_riak:get(privacy, {LUser, LServer}) of
+    case ejabberd_riak:get(privacy, privacy_schema(), {LUser, LServer}) of
         {ok, #privacy{default = Default, lists = Lists}} ->
             case Default of
                 none -> {none, []};
@@ -779,7 +785,7 @@ get_user_lists(LUser, LServer, mnesia) ->
             error
     end;
 get_user_lists(LUser, LServer, riak) ->
-    case ejabberd_riak:get(privacy, {LUser, LServer}) of
+    case ejabberd_riak:get(privacy, privacy_schema(), {LUser, LServer}) of
         {ok, #privacy{} = P} ->
             {ok, P};
         {error, _} ->
@@ -1242,6 +1248,6 @@ import(LServer) ->
 import(_LServer, mnesia, #privacy{} = P) ->
     mnesia:dirty_write(P);
 import(_LServer, riak, #privacy{} = P) ->
-    ejabberd_riak:put(P);
+    ejabberd_riak:put(P, privacy_schema());
 import(_, _, _) ->
     pass.