#!/usr/bin/python

import os

import sys

import random

import subprocess as sub

def identity(x):

return x

def cygpath(x):

command = ["cygpath", "-wp", x]

p = sub.Popen(command,stdout=sub.PIPE)

output, errors = p.communicate()

lines = output.split("\n")

return lines[0]

if sys.platform == "cygwin":

normclasspath = cygpath

else:

normclasspath = identity

CONF_DIR = os.path.expanduser("~/.storm")

STORM_DIR = "/".join(os.path.abspath( __file__ ).split("/")[:-2])

def get_jars_full(adir):

files = os.listdir(adir)

ret = []

for f in files:

if f.endswith(".jar"):

ret.append(adir + "/" + f)

return ret

def get_classpath(extrajars):

ret = get_jars_full(STORM_DIR)

ret.extend(get_jars_full(STORM_DIR + "/lib"))

ret.extend(extrajars)

return normclasspath(":".join(ret))

def confvalue(name, extrapaths):

cp = get_classpath(extrapaths)

command = ["java", "-client", "-cp", cp, "backtype.storm.command.config_value", name]

p = sub.Popen(command,stdout=sub.PIPE)

output, errors = p.communicate()

lines = output.split("\n")

for line in lines:

tokens = line.split(" ")

if tokens[0] == "VALUE:":

return tokens[1]

def print_localconfvalue(name):

print name + ": " + confvalue(name, [CONF_DIR])

def print_remoteconfvalue(name):

print name + ": " + confvalue(name, [STORM_DIR + "/conf"])

def exec_storm_class(klass, jvmtype="-server", childopts="", extrajars=[], args=[], prefix=""):

nativepath = confvalue("java.library.path", extrajars)

command = prefix + " java " + jvmtype + " -Djava.library.path=" + nativepath + " " + childopts + " -cp " + get_classpath(extrajars) + " " + klass + " " + " ".join(args)

print "Running: " + command

os.system(command)

def jar(jarfile, klass, *args):

exec_storm_class(

klass,

jvmtype="-client",

extrajars=[jarfile, CONF_DIR, STORM_DIR + "/bin"],

args=args,

prefix="export STORM_JAR=" + jarfile + ";")

def kill(*args):

exec_storm_class("backtype.storm.command.kill_topology", args=args, jvmtype="-client", extrajars=[CONF_DIR, STORM_DIR + "/bin"])

def shell(resourcesdir, command, *args):

tmpjarpath = "stormshell" + str(random.randint(0, 10000000)) + ".jar"

os.system("jar cf %s %s" % (tmpjarpath, resourcesdir))

runnerargs = [tmpjarpath, command]

runnerargs.extend(args)

exec_storm_class("backtype.storm.command.shell_submission", args=runnerargs, jvmtype="-client", extrajars=[CONF_DIR])

os.system("rm " + tmpjarpath)

def repl():

cppaths = [STORM_DIR + "/conf"]

exec_storm_class("clojure.lang.Repl", jvmtype="-client", extrajars=cppaths)

def nimbus():

cppaths = [STORM_DIR + "/log4j", STORM_DIR + "/conf"]

childopts = confvalue("nimbus.childopts", cppaths) + " -Dlogfile.name=nimbus.log -Dlog4j.configuration=storm.log.properties"

exec_storm_class("backtype.storm.daemon.nimbus", jvmtype="-server", extrajars=cppaths, childopts=childopts)

def supervisor():

cppaths = [STORM_DIR + "/log4j", STORM_DIR + "/conf"]

childopts = confvalue("nimbus.childopts", cppaths) + " -Dlogfile.name=supervisor.log -Dlog4j.configuration=storm.log.properties"

exec_storm_class("backtype.storm.daemon.supervisor", jvmtype="-server", extrajars=cppaths, childopts=childopts)

def ui():

childopts = "-Xmx768m -Dlogfile.name=ui.log -Dlog4j.configuration=storm.log.properties"

exec_storm_class("backtype.storm.ui.core", jvmtype="-server", childopts=childopts, extrajars=[STORM_DIR + "/log4j", STORM_DIR, STORM_DIR + "/conf"])

def drpc():

childopts = "-Xmx768m -Dlogfile.name=drpc.log -Dlog4j.configuration=storm.log.properties"

exec_storm_class("backtype.storm.daemon.drpc", jvmtype="-server", childopts=childopts, extrajars=[STORM_DIR + "/log4j", STORM_DIR + "/conf"])

def print_classpath():

print get_classpath([])

COMMANDS = {"jar": jar, "kill": kill, "shell": shell, "nimbus": nimbus, "ui": ui, "drpc": drpc, "supervisor": supervisor, "localconfvalue": print_localconfvalue, "remoteconfvalue": print_remoteconfvalue, "repl": repl, "classpath": print_classpath}

def print_commands():

global COMMANDS

cmds = COMMANDS.keys()

cmds.sort()

print "Commands:\n\t", reduce(lambda x,y: x + ', ' + y, cmds[1:], cmds[0])

def print_usage(msg=None):

if msg != None:

print msg

print_commands()

def main():

if len(sys.argv) <= 1:

print_usage()

sys.exit(-1)

COMMAND = sys.argv[1]

ARGS = sys.argv[2:]

COMMANDS[COMMAND](*ARGS)

if __name__ == "__main__":

main()