contrib/fb303/scripts/fb303_simple_mgmt.py

#!/usr/bin/env python
#
# Copyright (c) 2006- Facebook
# Distributed under the Thrift Software License
#
# See accompanying file LICENSE or visit the Thrift site at:
# http://developers.facebook.com/thrift/

import sys, os
from optparse import OptionParser

from thrift.Thrift import *

from thrift.transport import TSocket
from thrift.transport import TTransport
from thrift.protocol import TBinaryProtocol

from fb303 import *
from fb303.ttypes import *

def service_ctrl(
                 command,
                 port,
                 trans_factory = None,
                 prot_factory = None):
    """
    service_ctrl is a generic function to execute standard fb303 functions

    @param command: one of stop, start, reload, status, counters, name, alive
    @param port: service's port
    @param trans_factory: TTransportFactory to use for obtaining a TTransport. Default is
                          TBufferedTransportFactory
    @param prot_factory: TProtocolFactory to use for obtaining a TProtocol. Default is
                         TBinaryProtocolFactory
    """

    # Only root should be able to run these scripts, although we could relax this for some of the operations.
    if os.getuid() != 0:
        print "requires root."
        return 4

    if command in ["status"]:
        try:
            status = fb303_wrapper('status', port, trans_factory, prot_factory)
            status_details = fb303_wrapper('get_status_details', port, trans_factory, prot_factory)

            msg = fb_status_string(status)
            if (len(status_details)):
                msg += " - %s" % status_details
            print msg

            if (status == fb_status.ALIVE):
                return 2
            else:
                return 3
        except:
            print "Failed to get status"
            return 3

    # async commands
    if command in ["stop","reload"]:
        try:
            fb303_wrapper(command, port, trans_factory, prot_factory)
            return 0
        except:
            print "failed to tell the service to ", command
            return 3

    # scalar commands
    if command in ["version","alive","name"]:
        try:
            result = fb303_wrapper(command,  port, trans_factory, prot_factory)
            print result
            return 0
        except:
            print "failed to get ",command
            return 3

    # counters
    if command in ["counters"]:
        try:
            counters = fb303_wrapper('counters',  port, trans_factory, prot_factory)
            for counter in counters:
                print "%s: %d" % (counter, counters[counter])
            return 0
        except:
            print "failed to get counters"
            return 3

    return 0;


def fb303_wrapper(command, port, trans_factory = None, prot_factory = None):
    sock = TSocket.TSocket('localhost', port)

    # use input transport factory if provided
    if (trans_factory is None):
        trans = TTransport.TBufferedTransport(sock)
    else:
        trans = trans_factory.getTransport(sock)

    # use input protocol factory if provided
    if (prot_factory is None):
        prot = TBinaryProtocol.TBinaryProtocol(trans)
    else:
        prot = prot_factory.getProtocol(trans)

    # initialize client and open transport
    fb303_client = FacebookService.Client(prot, prot)
    trans.open()

    if (command == 'reload'):
        fb303_client.reinitialize()

    elif (command == 'stop'):
        fb303_client.shutdown()

    elif (command == 'status'):
        return fb303_client.getStatus()

    elif (command == 'version'):
        return fb303_client.getVersion()

    elif (command == 'get_status_details'):
        return fb303_client.getStatusDetails()

    elif (command == 'counters'):
        return fb303_client.getCounters()

    elif (command == 'name'):
        return fb303_client.getName()

    elif (command == 'alive'):
        return fb303_client.aliveSince()

    trans.close()


def fb_status_string(status_enum):
    if (status_enum == fb_status.DEAD):
        return "DEAD"
    if (status_enum == fb_status.STARTING):
        return "STARTING"
    if (status_enum == fb_status.ALIVE):
        return "ALIVE"
    if (status_enum == fb_status.STOPPING):
        return "STOPPING"
    if (status_enum == fb_status.STOPPED):
        return "STOPPED"
    if (status_enum == fb_status.WARNING):
        return "WARNING"


def main(port, command):
    status = service_ctrl(options.command, options.port)
    sys.exit(status)


# parse command line options
parser = OptionParser()
parser.add_option("-c", "--command", dest="command", help="execute this API", choices=["stop","counters","status","reload","version","name","alive"],
                  default="status")
parser.add_option("-p","--port",dest="port",help="the service's port", default=9082)

(options, args) = parser.parse_args()

main(options.port, options.command)