Skip to content

Commit 46805f4

Browse files
dumbbelldumbbell
committed
wip
1 parent ed44282 commit 46805f4

File tree

3 files changed

+105
-87
lines changed

3 files changed

+105
-87
lines changed

‎deps/rabbit/src/rabbit_cli.erl‎

Lines changed: 75 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -5,35 +5,56 @@
55
-export([main/1]).
66

77
main(Args) ->
8-
case run_cli(Args) of
9-
ok ->
10-
erlang:halt();
11-
{error, ErrorMsg} ->
12-
io:format(standard_error, "Error: ~ts~n", [ErrorMsg]),
13-
erlang:halt(1)
14-
end.
8+
Ret = run_cli(Args),
9+
io:format(standard_error, "Ret: ~p~n", [Ret]),
10+
erlang:halt().
1511

1612
run_cli(Args) ->
1713
maybe
1814
Progname = escript:script_name(),
1915
add_rabbitmq_code_path(Progname),
2016

21-
{ok, ArgMap, RemainingArgs} ?= parse_args(Progname, Args),
22-
Nodename = lookup_rabbitmq_nodename(ArgMap),
23-
{ok, _} ?= net_kernel:start(undefined, #{name_domain => shortnames}),
24-
25-
{ok, IO} ?= rabbit_cli_io:start_link(),
26-
Ret = case net_kernel:connect_node(Nodename) of
27-
true ->
28-
catch run_command(
29-
Nodename, Progname, ArgMap, RemainingArgs, IO);
30-
false ->
31-
catch run_command(
32-
undefined, Progname, ArgMap, RemainingArgs, IO)
33-
end,
34-
io:format("Ret = ~p~n", [Ret]),
35-
rabbit_cli_io:stop(IO),
36-
ok
17+
PartialArgparseDef = argparse_def(),
18+
{ok,
19+
PartialArgMap,
20+
PartialCmdPath,
21+
PartialCommand} ?= initial_parse(Progname, Args, PartialArgparseDef),
22+
23+
%% Get remote node name and prepare Erlang distribution.
24+
Nodename = lookup_rabbitmq_nodename(PartialArgMap),
25+
{ok, _} ?= net_kernel:start(
26+
undefined, #{name_domain => shortnames}),
27+
28+
{ok, IO} ?= rabbit_cli_io:start_link(Progname),
29+
try
30+
%% Can we reach the remote node?
31+
case net_kernel:connect_node(Nodename) of
32+
true ->
33+
maybe
34+
%% We can query the argparse definition from the
35+
%% remote node to know the commands it supports and
36+
%% proceed with the execution.
37+
ArgparseDef = get_final_argparse_def(Nodename),
38+
{ok,
39+
ArgMap,
40+
CmdPath,
41+
Command} ?= final_parse(Progname, Args, ArgparseDef),
42+
run_command(
43+
Nodename, ArgparseDef,
44+
Progname, ArgMap, CmdPath, Command,
45+
IO)
46+
end;
47+
false ->
48+
%% We can't reach the remote node. Let's fallback
49+
%% to a local execution.
50+
run_command(
51+
undefined, PartialArgparseDef,
52+
Progname, PartialArgMap, PartialCmdPath,
53+
PartialCommand, IO)
54+
end
55+
after
56+
rabbit_cli_io:stop(IO)
57+
end
3758
end.
3859

3960
add_rabbitmq_code_path(Progname) ->
@@ -73,34 +94,45 @@ argparse_def() ->
7394
short => $V,
7495
help =>
7596
"Display version and exit"}
76-
]}.
97+
],
98+
99+
commands => #{}}.
77100

78-
parse_args(Progname, Args) ->
79-
Definition = argparse_def(),
101+
initial_parse(Progname, Args, ArgparseDef) ->
80102
Options = #{progname => Progname},
81-
case partial_parse(Args, Definition, Options) of
82-
{ok, ArgMap, _CmdPath, _Command, RemainingArgs} ->
83-
{ok, ArgMap, RemainingArgs};
103+
case partial_parse(Args, ArgparseDef, Options) of
104+
{ok, ArgMap, CmdPath, Command, _RemainingArgs} ->
105+
{ok, ArgMap, CmdPath, Command};
84106
{error, _} = Error->
85107
Error
86108
end.
87109

88-
partial_parse(Args, Definition, Options) ->
89-
partial_parse(Args, Definition, Options, []).
110+
partial_parse(Args, ArgparseDef, Options) ->
111+
partial_parse(Args, ArgparseDef, Options, []).
90112

91-
partial_parse(Args, Definition, Options, RemainingArgs) ->
92-
case argparse:parse(Args, Definition, Options) of
113+
partial_parse(Args, ArgparseDef, Options, RemainingArgs) ->
114+
case argparse:parse(Args, ArgparseDef, Options) of
93115
{ok, ArgMap, CmdPath, Command} ->
94116
RemainingArgs1 = lists:reverse(RemainingArgs),
95117
{ok, ArgMap, CmdPath, Command, RemainingArgs1};
96118
{error, {_CmdPath, undefined, Arg, <<>>}} ->
97119
Args1 = Args -- [Arg],
98120
RemainingArgs1 = [Arg | RemainingArgs],
99-
partial_parse(Args1, Definition, Options, RemainingArgs1);
121+
partial_parse(Args1, ArgparseDef, Options, RemainingArgs1);
100122
{error, _} = Error ->
101123
Error
102124
end.
103125

126+
get_final_argparse_def(Nodename) ->
127+
ArgparseDef1 = argparse_def(),
128+
ArgparseDef2 = erpc:call(Nodename, rabbit_cli_commands, argparse_def, []),
129+
ArgparseDef = maps:merge(ArgparseDef1, ArgparseDef2),
130+
ArgparseDef.
131+
132+
final_parse(Progname, Args, ArgparseDef) ->
133+
Options = #{progname => Progname},
134+
argparse:parse(Args, ArgparseDef, Options).
135+
104136
lookup_rabbitmq_nodename(#{node := Nodename}) ->
105137
Nodename1 = complete_nodename(Nodename),
106138
Nodename1;
@@ -138,6 +170,15 @@ complete_nodename(Nodename) ->
138170
list_to_atom(Nodename)
139171
end.
140172

173+
run_command(
174+
_Nodename, ArgparseDef, _Progname, #{help := true}, CmdPath, _Command, IO) ->
175+
rabbit_cli_io:display_help(IO, CmdPath, ArgparseDef);
176+
run_command(Nodename, _ArgparseDef, Progname, ArgMap, CmdPath, Command, IO) ->
177+
erpc:call(
178+
Nodename,
179+
rabbit_cli_commands, run_command,
180+
[Progname, ArgMap, CmdPath, Command, IO]).
181+
141182
%lookup_command_map(Nodename) ->
142183
% %% Order of operations:
143184
% %% 1. refresh the cached copy:
@@ -178,17 +219,3 @@ complete_nodename(Nodename) ->
178219
% nomatch ->
179220
% true
180221
% end.
181-
182-
run_command(undefined, Progname, #{help := true}, _RemainingArgs, IO) ->
183-
Definition = argparse_def(),
184-
rabbit_cli_io:display_help(IO, Progname, [], Definition);
185-
run_command(Nodename, Progname, ArgMap, RemainingArgs, IO) ->
186-
try
187-
erpc:call(
188-
Nodename,
189-
rabbit_cli_commands, run_command,
190-
[Progname, ArgMap, RemainingArgs, IO])
191-
catch
192-
error:{erpc, Reason} ->
193-
{error, Reason}
194-
end.

‎deps/rabbit/src/rabbit_cli_commands.erl‎

Lines changed: 12 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -2,52 +2,40 @@
22

33
-include_lib("kernel/include/logger.hrl").
44

5-
-export([argparse_def/0, run_command/4]).
6-
-export([list_exchanges/3]).
5+
-export([argparse_def/0, run_command/5]).
6+
-export([list_exchanges/5]).
77

88
argparse_def() ->
99
%% Extract the commands from module attributes like feature flags and boot
1010
%% steps.
1111
#{commands =>
1212
#{"list" =>
13-
#{commands =>
13+
#{help => "List entities",
14+
commands =>
1415
#{"exchanges" =>
1516
maps:merge(
1617
rabbit_cli_io:argparse_def(record_stream),
17-
#{handler => {?MODULE, list_exchanges}})
18+
#{help => "List exchanges",
19+
handler => {?MODULE, list_exchanges}})
1820
}
1921
}
2022
}
2123
}.
2224

23-
run_command(Progname, ArgMap, Args, IO) ->
24-
Definition = argparse_def(),
25-
Options = #{progname => Progname},
26-
case argparse:parse(Args, Definition, Options) of
27-
{ok, NewArgMap, CmdPath, Command} ->
28-
ArgMap1 = maps:merge(ArgMap, NewArgMap),
29-
run_command1(Progname, CmdPath, ArgMap1, Command, IO);
30-
{error, Reason} = Error ->
31-
?LOG_ALERT("Error: ~s", [argparse:format_error(Reason)]),
32-
Error
33-
end.
34-
35-
run_command1(Progname, CmdPath, #{help := true}, Command, IO) ->
36-
rabbit_cli_io:display_help(IO, Progname, CmdPath, Command);
37-
run_command1(Progname, _CmdPath, ArgMap, Command, IO) ->
25+
run_command(Progname, ArgMap, CmdPath, Command, IO) ->
3826
%% TODO: Put both processes under the rabbit supervision tree.
39-
RunnerPid = command_runner(Progname, Command, ArgMap, IO),
27+
RunnerPid = command_runner(Progname, ArgMap, CmdPath, Command, IO),
4028
RunnerMRef = erlang:monitor(process, RunnerPid),
4129
receive
4230
{'DOWN', RunnerMRef, _, _, Reason} ->
4331
{ok, Reason}
4432
end.
4533

4634
command_runner(
47-
Progname, #{handler := {Mod, Fun}} = _Command, ArgMap, IO) ->
48-
spawn_link(Mod, Fun, [Progname, ArgMap, IO]).
35+
Progname, ArgMap, CmdPath, #{handler := {Mod, Fun}} = Command, IO) ->
36+
spawn_link(Mod, Fun, [Progname, ArgMap, CmdPath, Command, IO]).
4937

50-
list_exchanges(Progname, ArgMap, IO) ->
38+
list_exchanges(_Progname, ArgMap, _CmdPath, _Command, IO) ->
5139
InfoKeys = rabbit_exchange:info_keys(),
5240
Fields = lists:map(
5341
fun
@@ -70,7 +58,7 @@ list_exchanges(Progname, ArgMap, IO) ->
7058
(Key) ->
7159
#{name => Key, type => term}
7260
end, InfoKeys),
73-
case rabbit_cli_io:start_record_stream(IO, exchanges, Fields, {Progname, ArgMap}) of
61+
case rabbit_cli_io:start_record_stream(IO, exchanges, Fields, ArgMap) of
7462
{ok, Stream} ->
7563
Exchanges = rabbit_exchange:list(),
7664
lists:foreach(

‎deps/rabbit/src/rabbit_cli_io.erl‎

Lines changed: 18 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,10 @@
44

55
-include_lib("rabbit_common/include/resource.hrl").
66

7-
-export([start_link/0,
7+
-export([start_link/1,
88
stop/1,
99
argparse_def/1,
10-
display_help/4,
10+
display_help/3,
1111
start_record_stream/4,
1212
push_new_record/3,
1313
end_record_stream/2]).
@@ -18,10 +18,11 @@
1818
terminate/2,
1919
code_change/3]).
2020

21-
-record(?MODULE, {record_streams = #{}}).
21+
-record(?MODULE, {progname,
22+
record_streams = #{}}).
2223

23-
start_link() ->
24-
gen_server:start_link(rabbit_cli_io, none, []).
24+
start_link(Progname) ->
25+
gen_server:start_link(rabbit_cli_io, #{progname => Progname}, []).
2526

2627
stop(IO) ->
2728
MRef = erlang:monitor(process, IO),
@@ -49,28 +50,28 @@ argparse_def(record_stream) ->
4950
]
5051
}.
5152

52-
display_help(IO, Progname, CmdPath, ArgparseDef) ->
53-
gen_server:cast(IO, {?FUNCTION_NAME, Progname, CmdPath, ArgparseDef}).
53+
display_help(IO, CmdPath, Command) ->
54+
gen_server:cast(IO, {?FUNCTION_NAME, CmdPath, Command}).
5455

55-
start_record_stream(IO, Name, Fields, {_Progname, ArgMap} = ProgAndArgs)
56+
start_record_stream(IO, Name, Fields, ArgMap)
5657
when is_pid(IO) andalso
5758
is_atom(Name) andalso
5859
is_map(ArgMap) ->
59-
gen_server:call(IO, {?FUNCTION_NAME, Name, Fields, ProgAndArgs}).
60+
gen_server:call(IO, {?FUNCTION_NAME, Name, Fields, ArgMap}).
6061

6162
push_new_record(IO, #{name := Name}, Record) ->
6263
gen_server:cast(IO, {?FUNCTION_NAME, Name, Record}).
6364

6465
end_record_stream(IO, #{name := Name}) ->
6566
gen_server:cast(IO, {?FUNCTION_NAME, Name}).
6667

67-
init(_Args) ->
68+
init(#{progname := Progname}) ->
6869
process_flag(trap_exit, true),
69-
State = #?MODULE{},
70+
State = #?MODULE{progname = Progname},
7071
{ok, State}.
7172

7273
handle_call(
73-
{start_record_stream, Name, Fields, {_Progname, _ArgMap}},
74+
{start_record_stream, Name, Fields, _ArgMap},
7475
From,
7576
#?MODULE{record_streams = Streams} = State) ->
7677
Stream = #{name => Name, fields => Fields},
@@ -90,10 +91,12 @@ handle_call(_Request, _From, State) ->
9091
{reply, ok, State}.
9192

9293
handle_cast(
93-
{display_help, Progname, CmdPath, ArgparseDef},
94-
State) ->
94+
{display_help, CmdPath, ArgparseDef},
95+
#?MODULE{progname = Progname} = State) ->
9596
Options = #{progname => Progname,
96-
command => CmdPath},
97+
%% Work around bug in argparse;
98+
%% See https://github.com/erlang/otp/pull/9160
99+
command => tl(CmdPath)},
97100
Help = argparse:help(ArgparseDef, Options),
98101
io:format("~s~n", [Help]),
99102
{noreply, State};

0 commit comments

Comments
 (0)