]>
Commit | Line | Data |
---|---|---|
f67539c2 TL |
1 | #!/usr/bin/env python |
2 | ||
3 | # | |
4 | # Licensed to the Apache Software Foundation (ASF) under one | |
5 | # or more contributor license agreements. See the NOTICE file | |
6 | # distributed with this work for additional information | |
7 | # regarding copyright ownership. The ASF licenses this file | |
8 | # to you under the Apache License, Version 2.0 (the | |
9 | # "License"); you may not use this file except in compliance | |
10 | # with the License. You may obtain a copy of the License at | |
11 | # | |
12 | # http://www.apache.org/licenses/LICENSE-2.0 | |
13 | # | |
14 | # Unless required by applicable law or agreed to in writing, | |
15 | # software distributed under the License is distributed on an | |
16 | # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | |
17 | # KIND, either express or implied. See the License for the | |
18 | # specific language governing permissions and limitations | |
19 | # under the License. | |
20 | # | |
21 | ||
22 | from __future__ import print_function | |
23 | import sys | |
24 | import os | |
25 | from optparse import OptionParser | |
26 | ||
27 | from thrift.Thrift import * | |
28 | ||
29 | from thrift.transport import TSocket | |
30 | from thrift.transport import TTransport | |
31 | from thrift.protocol import TBinaryProtocol | |
32 | ||
33 | from fb303 import * | |
34 | from fb303.ttypes import * | |
35 | ||
36 | ||
37 | def service_ctrl( | |
38 | command, | |
39 | port, | |
40 | trans_factory=None, | |
41 | prot_factory=None): | |
42 | """ | |
43 | service_ctrl is a generic function to execute standard fb303 functions | |
44 | ||
45 | @param command: one of stop, start, reload, status, counters, name, alive | |
46 | @param port: service's port | |
47 | @param trans_factory: TTransportFactory to use for obtaining a TTransport. Default is | |
48 | TBufferedTransportFactory | |
49 | @param prot_factory: TProtocolFactory to use for obtaining a TProtocol. Default is | |
50 | TBinaryProtocolFactory | |
51 | """ | |
52 | ||
53 | if command in ["status"]: | |
54 | try: | |
55 | status = fb303_wrapper('status', port, trans_factory, prot_factory) | |
56 | status_details = fb303_wrapper('get_status_details', port, trans_factory, prot_factory) | |
57 | ||
58 | msg = fb_status_string(status) | |
59 | if (len(status_details)): | |
60 | msg += " - %s" % status_details | |
61 | print(msg) | |
62 | return 2 if status == fb_status.ALIVE else 3 | |
63 | except: | |
64 | print("Failed to get status") | |
65 | return 3 | |
66 | ||
67 | # scalar commands | |
68 | if command in ["version", "alive", "name"]: | |
69 | try: | |
70 | result = fb303_wrapper(command, port, trans_factory, prot_factory) | |
71 | print(result) | |
72 | return 0 | |
73 | except: | |
74 | print("failed to get ", command) | |
75 | return 3 | |
76 | ||
77 | # counters | |
78 | if command in ["counters"]: | |
79 | try: | |
80 | counters = fb303_wrapper('counters', port, trans_factory, prot_factory) | |
81 | for counter in counters: | |
82 | print("%s: %d" % (counter.encode('utf-8'), counters[counter])) | |
83 | return 0 | |
84 | except: | |
85 | print("failed to get counters") | |
86 | return 3 | |
87 | ||
88 | # Only root should be able to run the following commands | |
89 | if os.getuid() == 0: | |
90 | # async commands | |
91 | if command in ["stop", "reload"]: | |
92 | try: | |
93 | fb303_wrapper(command, port, trans_factory, prot_factory) | |
94 | return 0 | |
95 | except: | |
96 | print("failed to tell the service to ", command) | |
97 | return 3 | |
98 | else: | |
99 | if command in ["stop", "reload"]: | |
100 | print("root privileges are required to stop or reload the service.") | |
101 | return 4 | |
102 | ||
103 | print("The following commands are available:") | |
104 | for command in ["counters", "name", "version", "alive", "status"]: | |
105 | print("\t%s" % command) | |
106 | print("The following commands are available for users with root privileges:") | |
107 | for command in ["stop", "reload"]: | |
108 | print("\t%s" % command) | |
109 | ||
110 | return 0 | |
111 | ||
112 | ||
113 | def fb303_wrapper(command, port, trans_factory=None, prot_factory=None): | |
114 | sock = TSocket.TSocket('localhost', port) | |
115 | ||
116 | # use input transport factory if provided | |
117 | if (trans_factory is None): | |
118 | trans = TTransport.TBufferedTransport(sock) | |
119 | else: | |
120 | trans = trans_factory.getTransport(sock) | |
121 | ||
122 | # use input protocol factory if provided | |
123 | if (prot_factory is None): | |
124 | prot = TBinaryProtocol.TBinaryProtocol(trans) | |
125 | else: | |
126 | prot = prot_factory.getProtocol(trans) | |
127 | ||
128 | # initialize client and open transport | |
129 | fb303_client = FacebookService.Client(prot, prot) | |
130 | trans.open() | |
131 | ||
132 | if (command == 'reload'): | |
133 | fb303_client.reinitialize() | |
134 | ||
135 | elif (command == 'stop'): | |
136 | fb303_client.shutdown() | |
137 | ||
138 | elif (command == 'status'): | |
139 | return fb303_client.getStatus() | |
140 | ||
141 | elif (command == 'version'): | |
142 | return fb303_client.getVersion() | |
143 | ||
144 | elif (command == 'get_status_details'): | |
145 | return fb303_client.getStatusDetails() | |
146 | ||
147 | elif (command == 'counters'): | |
148 | return fb303_client.getCounters() | |
149 | ||
150 | elif (command == 'name'): | |
151 | return fb303_client.getName() | |
152 | ||
153 | elif (command == 'alive'): | |
154 | return fb303_client.aliveSince() | |
155 | ||
156 | trans.close() | |
157 | ||
158 | ||
159 | def fb_status_string(status_enum): | |
160 | if (status_enum == fb_status.DEAD): | |
161 | return "DEAD" | |
162 | if (status_enum == fb_status.STARTING): | |
163 | return "STARTING" | |
164 | if (status_enum == fb_status.ALIVE): | |
165 | return "ALIVE" | |
166 | if (status_enum == fb_status.STOPPING): | |
167 | return "STOPPING" | |
168 | if (status_enum == fb_status.STOPPED): | |
169 | return "STOPPED" | |
170 | if (status_enum == fb_status.WARNING): | |
171 | return "WARNING" | |
172 | ||
173 | ||
174 | def main(): | |
175 | ||
176 | # parse command line options | |
177 | parser = OptionParser() | |
178 | commands = ["stop", "counters", "status", "reload", "version", "name", "alive"] | |
179 | ||
180 | parser.add_option("-c", "--command", dest="command", help="execute this API", | |
181 | choices=commands, default="status") | |
182 | parser.add_option("-p", "--port", dest="port", help="the service's port", | |
183 | default=9082) | |
184 | ||
185 | (options, args) = parser.parse_args() | |
186 | status = service_ctrl(options.command, options.port) | |
187 | sys.exit(status) | |
188 | ||
189 | ||
190 | if __name__ == '__main__': | |
191 | main() |