Merge pull request #241 from twitter/karthik/ps

added ps command

heron/cli/src/python/BUILD

@@ -13,6 +13,7 @@ pex_library(
         "jars.py",
         "kill.py",
         "opts.py",
+        "ps.py",
         "restart.py",
         "submit.py",
         "utils.py",

heron/cli/src/python/jars.py

@@ -10,14 +10,23 @@ def pick(dirname, pattern):
   return file_list[0] if file_list else None
 
 ################################################################################
-# Get the topology jars - TODO, make the jars independent version free
+# Get the topology jars
 ################################################################################
 def topology_jars():
   jars = [
       os.path.join(utils.get_heron_lib_dir(), "3rdparty", "*")
   ]
   return jars
 
+################################################################################
+# Get the command jars
+################################################################################
+def command_jars():
+  jars = [
+      os.path.join(utils.get_heron_lib_dir(), "commands", "*")
+  ]
+  return jars
+
 ################################################################################
 # Get the scheduler jars
 ################################################################################

heron/cli/src/python/main.py

@@ -23,6 +23,7 @@
 import heron.cli.src.python.activate as activate
 import heron.cli.src.python.deactivate as deactivate
 import heron.cli.src.python.kill as kill
+import heron.cli.src.python.ps as ps
 import heron.cli.src.python.restart as restart
 import heron.cli.src.python.submit as submit
 import heron.cli.src.python.utils as utils
@@ -71,6 +72,7 @@ def create_parser():
   deactivate.create_parser(subparsers)
   help.create_parser(subparsers)
   kill.create_parser(subparsers)
+  ps.create_parser(subparsers)
   restart.create_parser(subparsers)
   submit.create_parser(subparsers)
   version.create_parser(subparsers)
@@ -93,6 +95,9 @@ def run(command, parser, command_args, unknown_args):
   elif command == 'restart':
     return restart.run(command, parser, command_args, unknown_args)
 
+  elif command == 'ps':
+    return ps.run(command, parser, command_args, unknown_args)
+
   elif command == 'submit':
     return submit.run(command, parser, command_args, unknown_args)
 

heron/cli/src/python/ps.py

@@ -0,0 +1,70 @@
+#!/usr/bin/python2.7
+
+import argparse
+import atexit
+import base64
+import contextlib
+import glob
+import logging
+import logging.handlers
+import os
+import shutil
+import sys
+import subprocess
+import tarfile
+import tempfile
+
+from heron.common.src.python.color import Log
+
+import heron.cli.src.python.args as args
+import heron.cli.src.python.execute as execute
+import heron.cli.src.python.jars as jars
+import heron.cli.src.python.utils as utils
+
+def create_parser(subparsers):
+  parser = subparsers.add_parser(
+      'ps',
+      help='List all topologies',
+      usage = "%(prog)s [options] cluster/[role]/[environ]",
+      add_help = False)
+
+  args.add_titles(parser)
+  args.add_cluster_role_env(parser)
+
+  args.add_config(parser)
+  args.add_verbose(parser)
+
+  parser.set_defaults(subcommand='ps')
+  return parser
+
+def run(command, parser, cl_args, unknown_args):
+
+  try:
+    config_overrides = utils.parse_cmdline_override(cl_args)
+
+    new_args = [
+        "--cluster", cl_args['cluster'],
+        "--role", cl_args['role'],
+        "--environment", cl_args['environ'],
+        "--heron_home", utils.get_heron_dir(),
+        "--config_path", cl_args['config_path'],
+        "--config_overrides", base64.b64encode(config_overrides),
+        "--command", "ps",
+    ]
+
+    lib_jars = utils.get_heron_libs(jars.command_jars() + jars.statemgr_jars())
+
+    # invoke the runtime manager to kill the topology
+    execute.heron_class(
+        'com.twitter.heron.command.CommandHandlerMain',
+        lib_jars,
+        extra_jars=[],
+        args= new_args
+    )
+
+  except Exception as ex:
+    print ex
+    Log.error('Failed to get list of topologies')
+    return False
+
+  return True

heron/commands/src/java/BUILD

@@ -0,0 +1,41 @@
+package(default_visibility = ["//visibility:public"])
+
+load("/tools/rules/heron_deps", "heron_java_proto_files")
+
+common_deps_files = [
+   "//heron/common/src/java:common-java",
+   "//3rdparty/commons:commons-cli-java",
+   "//3rdparty/guava:guava-java",
+]
+
+spi_deps_files = [
+    "//heron/spi/src/java:common-spi-java",
+    "//heron/spi/src/java:statemgr-spi-java",
+    "//heron/spi/src/java:utils-spi-java",
+]
+
+commands_deps_files = \
+    common_deps_files + \
+    heron_java_proto_files() + \
+    spi_deps_files
+
+java_library(
+    name = 'commands-java',
+    srcs = glob(
+        ["**/*.java"],
+    ),
+    deps = commands_deps_files,
+)
+
+java_binary(
+    name = 'commands-unshaded',
+    srcs = glob(["**/*.java"]),
+    deps = commands_deps_files,
+)
+
+genrule(
+    name = "heron-commands",
+    srcs = [":commands-unshaded_deploy.jar"],
+    outs = ["heron-commands.jar"],
+    cmd  = "cp $< $@",
+)

heron/commands/src/java/com/twitter/heron/command/CommandHandler.java

@@ -0,0 +1,40 @@
+package com.twitter.heron.command;
+
+import java.io.IOException;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+import com.twitter.heron.spi.common.Config;
+import com.twitter.heron.spi.common.Context;
+
+public abstract class CommandHandler {
+
+  // static config read from the config files
+  protected Config config;
+
+  // runtime config gathered during execution
+  protected Config runtime;
+
+  /**
+   * Construct the command handler with static and runtime config
+   */
+  CommandHandler(Config config, Config runtime) {
+    this.config = config;
+    this.runtime = runtime;
+  }
+
+  /**
+   * Execute any conditions before the command execution
+   */
+  public abstract boolean beforeExecution() throws Exception;
+
+  /**
+   * Execute any cleanup after the command execution
+   */
+  public abstract boolean afterExecution() throws Exception;
+
+  /** 
+   * Execute the command
+   */
+  public abstract boolean execute() throws Exception;
+}

heron/commands/src/java/com/twitter/heron/command/CommandHandlerConfig.java

@@ -0,0 +1,79 @@
+package com.twitter.heron.command;
+
+import java.io.IOException;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+import com.twitter.heron.api.generated.TopologyAPI;
+import com.twitter.heron.common.basics.FileUtils;
+import com.twitter.heron.proto.scheduler.Scheduler;
+
+import com.twitter.heron.spi.common.ClusterConfig;
+import com.twitter.heron.spi.common.ClusterDefaults;
+import com.twitter.heron.spi.common.Config;
+import com.twitter.heron.spi.common.Context;
+import com.twitter.heron.spi.common.Keys;
+
+import org.apache.commons.cli.CommandLine;
+
+/**
+ * For loading command handler config
+ */
+public class CommandHandlerConfig {
+  private static final Logger LOG = Logger.getLogger(CommandHandlerConfig.class.getName());
+
+  /**
+   * Load the defaults config
+   *
+   * @return config, the defaults config
+   */
+  protected static Config defaultConfigs(String heronHome, String configPath) {
+    Config config = Config.newBuilder()
+        .putAll(ClusterDefaults.getDefaults())
+        .putAll(ClusterConfig.loadCommandsConfig(heronHome, configPath))
+        .build();
+    return config;
+  }
+
+  /**
+   * Load the config parameters from the command line
+   *
+   * @param cluster, name of the cluster
+   * @param role, user role
+   * @param environ, user provided environment/tag
+   *
+   * @return config, the command line config
+   */
+  protected static Config commandLineConfigs(String cluster, String role, String environ) {
+    Config config = Config.newBuilder()
+        .put(Keys.cluster(), cluster)
+        .put(Keys.role(), role)
+        .put(Keys.environ(), environ)
+        .build();
+    return config;
+  }
+
+  /**
+   * Load the config from static config files
+   *
+   * @param commandLine, the command line args provided
+   *
+   * @return config, the static config
+   */
+  protected static Config loadConfig(CommandLine commandLine) {
+
+    String cluster = commandLine.getOptionValue("cluster");
+    String role = commandLine.getOptionValue("role");
+    String environ = commandLine.getOptionValue("environment");
+    String heronHome = commandLine.getOptionValue("heron_home");
+    String configPath = commandLine.getOptionValue("config_path");
+
+    // build the config by expanding all the variables
+    Config config = Config.expand(
+        Config.newBuilder()
+            .putAll(defaultConfigs(heronHome, configPath))
+            .putAll(commandLineConfigs(cluster, role, environ))
+            .build());
+    return config;
+  }
+}

heron/commands/src/java/com/twitter/heron/command/CommandHandlerFactory.java

@@ -0,0 +1,20 @@
+package com.twitter.heron.command;
+
+import java.io.IOException;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+import com.twitter.heron.spi.common.Config;
+import com.twitter.heron.spi.common.Context;
+
+public class CommandHandlerFactory {
+  private static final Logger LOG = Logger.getLogger(CommandHandlerFactory.class.getName());
+
+  public static CommandHandler makeCommand(String command, Config config, Config runtime) {
+    if (command.equalsIgnoreCase("ps"))
+      return new ListTopologiesHandler(config, runtime);
+
+    LOG.info("Invalid command " + command);
+    return null;
+  }
+}