浏览代码

[spark] Allow session backend to be picked on the cli

Erick Tryzelaar 11 年之前
父节点
当前提交
9821fffcc1
共有 1 个文件被更改,包括 30 次插入2 次删除
  1. 30 2
      apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/Main.scala

+ 30 - 2
apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/Main.scala

@@ -7,13 +7,35 @@ import org.scalatra._
 import org.scalatra.servlet.ScalatraListener
 
 object Main {
+
+  val SESSION_KIND = "livy.session.kind"
+  val PROCESS_SESSION = "process"
+  val YARN_SESSION = "yarn"
+
   def main(args: Array[String]): Unit = {
+
+    if (args.length != 1) {
+      println("Must specify either `process` or `yarn` for the session kind")
+      sys.exit(1)
+    }
+
+    val session_kind = args(0)
+
+    session_kind match {
+      case PROCESS_SESSION | YARN_SESSION =>
+      case _ =>
+        println("Unknown session kind: " + session_kind)
+        sys.exit(1)
+    }
+
     val port = sys.env.getOrElse("PORT", "8998").toInt
     val server = new WebServer(port)
 
+
     server.context.setResourceBase("src/main/com/cloudera/hue/livy/server")
     server.context.setInitParameter(ScalatraListener.LifeCycleKey, classOf[ScalatraBootstrap].getCanonicalName)
     server.context.addEventListener(new ScalatraListener)
+    server.context.setInitParameter(SESSION_KIND, session_kind)
 
     server.start()
     server.join()
@@ -23,11 +45,17 @@ object Main {
 
 class ScalatraBootstrap extends LifeCycle {
 
-  val sessionFactory = new YarnSessionFactory
-  val sessionManager = new SessionManager(sessionFactory)
+  var sessionManager: SessionManager = null
 
   override def init(context: ServletContext): Unit = {
     context.mount(new WebApp(sessionManager), "/*")
+
+    val sessionFactory = context.getInitParameter(Main.SESSION_KIND) match {
+      case Main.PROCESS_SESSION => new ProcessSessionFactory
+      case Main.YARN_SESSION => new YarnSessionFactory
+    }
+
+    sessionManager = new SessionManager(sessionFactory)
   }
 
   override def destroy(context: ServletContext): Unit = {