PageRenderTime 30ms CodeModel.GetById 16ms app.highlight 11ms RepoModel.GetById 1ms app.codeStats 0ms

/tags/release-0.1-rc2/hive/external/service/lib/py/fb303_scripts/fb303_simple_mgmt.py

#
Python | 195 lines | 164 code | 6 blank | 25 comment | 4 complexity | b525aa2860045856c3046fbea5023598 MD5 | raw file
  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
 22import sys, os
 23from optparse import OptionParser
 24
 25from thrift.Thrift import *
 26
 27from thrift.transport import TSocket
 28from thrift.transport import TTransport
 29from thrift.protocol import TBinaryProtocol
 30
 31from fb303 import *
 32from fb303.ttypes import *
 33
 34def service_ctrl(
 35                 command,
 36                 port,
 37                 trans_factory = None,
 38                 prot_factory = None):
 39    """
 40    service_ctrl is a generic function to execute standard fb303 functions
 41
 42    @param command: one of stop, start, reload, status, counters, name, alive
 43    @param port: service's port
 44    @param trans_factory: TTransportFactory to use for obtaining a TTransport. Default is
 45                          TBufferedTransportFactory
 46    @param prot_factory: TProtocolFactory to use for obtaining a TProtocol. Default is
 47                         TBinaryProtocolFactory
 48    """
 49
 50    if command in ["status"]:
 51        try:
 52            status = fb303_wrapper('status', port, trans_factory, prot_factory)
 53            status_details = fb303_wrapper('get_status_details', port, trans_factory, prot_factory)
 54
 55            msg = fb_status_string(status)
 56            if (len(status_details)):
 57                msg += " - %s" % status_details
 58            print msg
 59
 60            if (status == fb_status.ALIVE):
 61                return 2
 62            else:
 63                return 3
 64        except:
 65            print "Failed to get status"
 66            return 3
 67
 68    # scalar commands
 69    if command in ["version","alive","name"]:
 70        try:
 71            result = fb303_wrapper(command,  port, trans_factory, prot_factory)
 72            print result
 73            return 0
 74        except:
 75            print "failed to get ",command
 76            return 3
 77
 78    # counters
 79    if command in ["counters"]:
 80        try:
 81            counters = fb303_wrapper('counters',  port, trans_factory, prot_factory)
 82            for counter in counters:
 83                print "%s: %d" % (counter, counters[counter])
 84            return 0
 85        except:
 86            print "failed to get counters"
 87            return 3
 88
 89
 90    # Only root should be able to run the following commands
 91    if os.getuid() == 0:
 92        # async commands
 93        if command in ["stop","reload"] :
 94            try:
 95                fb303_wrapper(command, port, trans_factory, prot_factory)
 96                return 0
 97            except:
 98                print "failed to tell the service to ", command
 99                return 3
100    else:
101        if command in ["stop","reload"]:
102            print "root privileges are required to stop or reload the service."
103            return 4
104
105    print "The following commands are available:"
106    for command in ["counters","name","version","alive","status"]:
107        print "\t%s" % command
108    print "The following commands are available for users with root privileges:"
109    for command in ["stop","reload"]:
110        print "\t%s" % command
111
112
113
114    return 0;
115
116
117def fb303_wrapper(command, port, trans_factory = None, prot_factory = None):
118    sock = TSocket.TSocket('localhost', port)
119
120    # use input transport factory if provided
121    if (trans_factory is None):
122        trans = TTransport.TBufferedTransport(sock)
123    else:
124        trans = trans_factory.getTransport(sock)
125
126    # use input protocol factory if provided
127    if (prot_factory is None):
128        prot = TBinaryProtocol.TBinaryProtocol(trans)
129    else:
130        prot = prot_factory.getProtocol(trans)
131
132    # initialize client and open transport
133    fb303_client = FacebookService.Client(prot, prot)
134    trans.open()
135
136    if (command == 'reload'):
137        fb303_client.reinitialize()
138
139    elif (command == 'stop'):
140        fb303_client.shutdown()
141
142    elif (command == 'status'):
143        return fb303_client.getStatus()
144
145    elif (command == 'version'):
146        return fb303_client.getVersion()
147
148    elif (command == 'get_status_details'):
149        return fb303_client.getStatusDetails()
150
151    elif (command == 'counters'):
152        return fb303_client.getCounters()
153
154    elif (command == 'name'):
155        return fb303_client.getName()
156
157    elif (command == 'alive'):
158        return fb303_client.aliveSince()
159
160    trans.close()
161
162
163def fb_status_string(status_enum):
164    if (status_enum == fb_status.DEAD):
165        return "DEAD"
166    if (status_enum == fb_status.STARTING):
167        return "STARTING"
168    if (status_enum == fb_status.ALIVE):
169        return "ALIVE"
170    if (status_enum == fb_status.STOPPING):
171        return "STOPPING"
172    if (status_enum == fb_status.STOPPED):
173        return "STOPPED"
174    if (status_enum == fb_status.WARNING):
175        return "WARNING"
176
177
178def main():
179
180    # parse command line options
181    parser = OptionParser()
182    commands=["stop","counters","status","reload","version","name","alive"]
183
184    parser.add_option("-c", "--command", dest="command", help="execute this API",
185                      choices=commands, default="status")
186    parser.add_option("-p","--port",dest="port",help="the service's port",
187                      default=9082)
188
189    (options, args) = parser.parse_args()
190    status = service_ctrl(options.command, options.port)
191    sys.exit(status)
192
193
194if __name__ == '__main__':
195    main()