Subversion Repositories SmartDukaan

#!/usr/bin/env python

#

# Licensed to the Apache Software Foundation (ASF) under one

# or more contributor license agreements. See the NOTICE file

# distributed with this work for additional information

# regarding copyright ownership. The ASF licenses this file

# to you under the Apache License, Version 2.0 (the

# "License"); you may not use this file except in compliance

# with the License. You may obtain a copy of the License at

#

#   http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing,

# software distributed under the License is distributed on an

# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY

# KIND, either express or implied. See the License for the

# specific language governing permissions and limitations

# under the License.

#

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

    """

    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

    # 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

    # Only root should be able to run the following commands

    if os.getuid() == 0:

        # 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

    else:

        if command in ["stop","reload"]:

            print "root privileges are required to stop or reload the service."

            return 4

    print "The following commands are available:"

    for command in ["counters","name","version","alive","status"]:

        print "\t%s" % command

    print "The following commands are available for users with root privileges:"

    for command in ["stop","reload"]:

        print "\t%s" % command

    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():

    # parse command line options

    parser = OptionParser()

    commands=["stop","counters","status","reload","version","name","alive"]

    parser.add_option("-c", "--command", dest="command", help="execute this API",

                      choices=commands, default="status")

    parser.add_option("-p","--port",dest="port",help="the service's port",

                      default=9082)

    (options, args) = parser.parse_args()

    status = service_ctrl(options.command, options.port)

    sys.exit(status)

if __name__ == '__main__':

    main()