improve get_entity_* to work regardless of nodetree (EJAB-937)

SVN Revision: 2104

6 files changed, 53 insertions, 31 deletions

diff --git a/src/mod_pubsub/gen_pubsub_nodetree.erl b/src/mod_pubsub/gen_pubsub_nodetree.erl
index 941748e68..f012b3d0b 100644
--- a/src/mod_pubsub/gen_pubsub_nodetree.erl
+++ b/src/mod_pubsub/gen_pubsub_nodetree.erl
@@ -44,6 +44,7 @@ behaviour_info(callbacks) ->
 	 {set_node, 1},
 	 {get_node, 3},
 	 {get_node, 2},
+	 {get_node, 1},
 	 {get_nodes, 2},
 	 {get_nodes, 1},
 	 {get_subnodes, 3},
diff --git a/src/mod_pubsub/mod_pubsub.erl b/src/mod_pubsub/mod_pubsub.erl
index 4129fd628..e0acdd2a2 100644
--- a/src/mod_pubsub/mod_pubsub.erl
+++ b/src/mod_pubsub/mod_pubsub.erl
@@ -165,12 +165,12 @@ init([ServerHost, Opts]) ->
     PepOffline = gen_mod:get_opt(pep_sendlast_offline, Opts, false),
     IQDisc = gen_mod:get_opt(iqdisc, Opts, one_queue),
     pubsub_index:init(Host, ServerHost, Opts),
+    ets:new(gen_mod:get_module_proc(Host, pubsub_state), [set, named_table]),
+    ets:new(gen_mod:get_module_proc(ServerHost, pubsub_state), [set, named_table]),
     {Plugins, NodeTree, PepMapping} = init_plugins(Host, ServerHost, Opts),
     mod_disco:register_feature(ServerHost, ?NS_PUBSUB),
-    ets:new(gen_mod:get_module_proc(Host, pubsub_state), [set, named_table]),
     ets:insert(gen_mod:get_module_proc(Host, pubsub_state), {nodetree, NodeTree}),
     ets:insert(gen_mod:get_module_proc(Host, pubsub_state), {plugins, Plugins}),
-    ets:new(gen_mod:get_module_proc(ServerHost, pubsub_state), [set, named_table]),
     ets:insert(gen_mod:get_module_proc(ServerHost, pubsub_state), {nodetree, NodeTree}),
     ets:insert(gen_mod:get_module_proc(ServerHost, pubsub_state), {plugins, Plugins}),
     ets:insert(gen_mod:get_module_proc(ServerHost, pubsub_state), {pep_mapping, PepMapping}),
@@ -3003,6 +3003,7 @@ tree_call(Host, Function, Args) ->
     end,
     catch apply(Module, Function, Args).
 tree_action(Host, Function, Args) ->
+    ?DEBUG("tree_action ~p ~p ~p",[Host,Function,Args]),
     Fun = fun() -> tree_call(Host, Function, Args) end,
     catch mnesia:sync_dirty(Fun).
 
@@ -3022,7 +3023,8 @@ node_call(Type, Function, Args) ->
 	Result -> {result, Result} %% any other return value is forced as result
     end.
 
-node_action(_Host, Type, Function, Args) ->
+node_action(Host, Type, Function, Args) ->
+    ?DEBUG("node_action ~p ~p ~p ~p",[Host,Type,Function,Args]),
     transaction(fun() ->
 			node_call(Type, Function, Args)
 		end, sync_dirty).
diff --git a/src/mod_pubsub/node_default.erl b/src/mod_pubsub/node_default.erl
index ad7ad73bf..6e39b88e8 100644
--- a/src/mod_pubsub/node_default.erl
+++ b/src/mod_pubsub/node_default.erl
@@ -553,15 +553,14 @@ get_entity_affiliations(Host, Owner) ->
     SubKey = jlib:jid_tolower(Owner),
     GenKey = jlib:jid_remove_resource(SubKey),
     States = mnesia:match_object(#pubsub_state{stateid = {GenKey, '_'}, _ = '_'}),
+    NodeTree = case ets:lookup(gen_mod:get_module_proc(Host, pubsub_state), nodetree) of
+	    [{nodetree, N}] -> N;
+	    _ -> nodetree_default
+	end,
     Reply = lists:foldl(fun(#pubsub_state{stateid = {_, N}, affiliation = A}, Acc) ->
-	case mnesia:index_read(pubsub_node, N, #pubsub_node.id) of
-	    [#pubsub_node{nodeid = {H, _}} = Node] ->
-		case H of
-		    Host -> [{Node, A}|Acc];
-		    _ -> Acc
-		end;
-	    _ ->
-		Acc
+	case NodeTree:get_node(N) of
+	    #pubsub_node{nodeid = {Host, _}} = Node -> [{Node, A}|Acc];
+	    _ -> Acc
 	end
     end, [], States),
     {result, Reply}.
@@ -611,15 +610,14 @@ get_entity_subscriptions(Host, Owner) ->
 	    ++ mnesia:match_object(
 	       #pubsub_state{stateid = {SubKey, '_'}, _ = '_'})
     end,
+    NodeTree = case ets:lookup(gen_mod:get_module_proc(Host, pubsub_state), nodetree) of
+	    [{nodetree, N}] -> N;
+	    _ -> nodetree_default
+	end,
     Reply = lists:foldl(fun(#pubsub_state{stateid = {J, N}, subscription = S}, Acc) ->
-	case mnesia:index_read(pubsub_node, N, #pubsub_node.id) of
-	    [#pubsub_node{nodeid = {H, _}} = Node] ->
-		case H of
-		    Host -> [{Node, S, J}|Acc];
-		    _ -> Acc
-		end;
-	    _ ->
-		Acc
+	case NodeTree:get_node(N) of
+	    #pubsub_node{nodeid = {Host, _}} = Node -> [{Node, S, J}|Acc];
+	    _ -> Acc
 	end
     end, [], States),
     {result, Reply}.
diff --git a/src/mod_pubsub/node_pep.erl b/src/mod_pubsub/node_pep.erl
index 833612ae3..f0f7ac4b0 100644
--- a/src/mod_pubsub/node_pep.erl
+++ b/src/mod_pubsub/node_pep.erl
@@ -168,8 +168,21 @@ purge_node(NodeId, Owner) ->
     node_default:purge_node(NodeId, Owner).
 
 get_entity_affiliations(_Host, Owner) ->
-    OwnerKey = jlib:jid_tolower(jlib:jid_remove_resource(Owner)),
-    node_default:get_entity_affiliations(OwnerKey, Owner).
+    {_, D, _} = SubKey = jlib:jid_tolower(Owner),
+    SubKey = jlib:jid_tolower(Owner),
+    GenKey = jlib:jid_remove_resource(SubKey),
+    States = mnesia:match_object(#pubsub_state{stateid = {GenKey, '_'}, _ = '_'}),
+    NodeTree = case ets:lookup(gen_mod:get_module_proc(D, pubsub_state), nodetree) of
+	    [{nodetree, N}] -> N;
+	    _ -> nodetree_default
+	end,
+    Reply = lists:foldl(fun(#pubsub_state{stateid = {_, N}, affiliation = A}, Acc) ->
+	case NodeTree:get_node(N) of
+	    #pubsub_node{nodeid = {{_, D, _}, _}} = Node -> [{Node, A}|Acc];
+	    _ -> Acc
+	end
+    end, [], States),
+    {result, Reply}.
 
 get_node_affiliations(NodeId) ->
     node_default:get_node_affiliations(NodeId).
@@ -191,15 +204,14 @@ get_entity_subscriptions(_Host, Owner) ->
 	    ++ mnesia:match_object(
 	       #pubsub_state{stateid = {SubKey, '_'}, _ = '_'})
     end,
+    NodeTree = case ets:lookup(gen_mod:get_module_proc(D, pubsub_state), nodetree) of
+	    [{nodetree, N}] -> N;
+	    _ -> nodetree_default
+	end,
     Reply = lists:foldl(fun(#pubsub_state{stateid = {J, N}, subscription = S}, Acc) ->
-	case mnesia:index_read(pubsub_node, N, #pubsub_node.id) of
-	    [#pubsub_node{nodeid = {H, _}} = Node] ->
-		case H of
-		    {_, D, _} -> [{Node, S, J}|Acc];
-		    _ -> Acc
-		end;
-	    _ ->
-		Acc
+	case NodeTree:get_node(N) of
+	    #pubsub_node{nodeid = {{_, D, _}, _}} = Node -> [{Node, S, J}|Acc];
+	    _ -> Acc
 	end
     end, [], States),
     {result, Reply}.
@@ -265,4 +277,3 @@ complain_if_modcaps_disabled(ServerHost) ->
 	_ ->
 	    ok
     end.
-
diff --git a/src/mod_pubsub/nodetree_default.erl b/src/mod_pubsub/nodetree_default.erl
index 284fe7f8e..ac27cf3c5 100644
--- a/src/mod_pubsub/nodetree_default.erl
+++ b/src/mod_pubsub/nodetree_default.erl
@@ -47,6 +47,7 @@
 	 set_node/1,
 	 get_node/3,
 	 get_node/2,
+	 get_node/1,
 	 get_nodes/2,
 	 get_nodes/1,
 	 get_subnodes/3,
@@ -109,6 +110,12 @@ get_node(Host, Node) ->
 	[] -> {error, ?ERR_ITEM_NOT_FOUND};
 	Error -> Error
     end.
+get_node(NodeId) ->
+    case catch mnesia:index_read(pubsub_node, NodeId, #pubsub_node.id) of
+	[Record] when is_record(Record, pubsub_node) -> Record;
+	[] -> {error, ?ERR_ITEM_NOT_FOUND};
+	Error -> Error
+    end.
 
 %% @spec (Host) -> [pubsubNode()] | {error, Reason}
 %%     Host = mod_pubsub:host() | mod_pubsub:jid()
diff --git a/src/mod_pubsub/nodetree_virtual.erl b/src/mod_pubsub/nodetree_virtual.erl
index 05107cb7e..43b516072 100644
--- a/src/mod_pubsub/nodetree_virtual.erl
+++ b/src/mod_pubsub/nodetree_virtual.erl
@@ -44,6 +44,7 @@
 	 set_node/1,
 	 get_node/3,
 	 get_node/2,
+	 get_node/1,
 	 get_nodes/2,
 	 get_nodes/1,
 	 get_subnodes/3,
@@ -90,7 +91,9 @@ set_node(_NodeRecord) ->
 get_node(Host, Node, _From) ->
     get_node(Host, Node).
 get_node(Host, Node) ->
-    #pubsub_node{nodeid = {Host, Node}}.
+    #pubsub_node{nodeid = {Host, Node}, id = {Host, Node}}.
+get_node(NodeId) ->
+    #pubsub_node{nodeid = NodeId, id = NodeId}.
 
 %% @spec (Host) -> [pubsubNode()]
 %%     Host = mod_pubsub:host() | mod_pubsub:jid()