[livy] Turn session kind into a case class for added type safety

apps/spark/java/livy-repl/src/main/scala/com/cloudera/hue/livy/repl/Main.scala

@@ -5,7 +5,7 @@ import javax.servlet.ServletContext
 
 import com.cloudera.hue.livy.repl.python.PythonSession
 import com.cloudera.hue.livy.repl.scala.SparkSession
-import com.cloudera.hue.livy.{Logging, WebServer}
+import com.cloudera.hue.livy.{Utils, Logging, WebServer}
 import dispatch._
 import org.json4s.jackson.Serialization.write
 import org.json4s.{DefaultFormats, Formats}
@@ -40,10 +40,10 @@ object Main extends Logging {
       sys.exit(1)
     }
 
-    val session_kind = args(0)
+    val session_kind = args.head
 
     session_kind match {
-      case PYTHON_SESSION | PYSPARK_SESSION | SCALA_SESSION | SPARK_SESSION =>
+      case PYSPARK_SESSION | SPARK_SESSION =>
       case _ =>
         println("Unknown session kind: " + session_kind)
         sys.exit(1)
@@ -83,10 +83,8 @@ class ScalatraBootstrap extends LifeCycle with Logging {
 
   override def init(context: ServletContext): Unit = {
     session = context.getInitParameter(Main.SESSION_KIND) match {
-      case Main.PYTHON_SESSION => PythonSession.createPySpark()
-      case Main.PYSPARK_SESSION => PythonSession.createPySpark()
-      case Main.SCALA_SESSION => SparkSession.create()
-      case Main.SPARK_SESSION => SparkSession.create()
+      case Main.PYSPARK_SESSION | Main.PYTHON_SESSION => PythonSession.createPySpark()
+      case Main.SPARK_SESSION | Main.SCALA_SESSION => SparkSession.create()
     }
 
     context.mount(new WebApp(session), "/*")

apps/spark/java/livy-repl/src/main/scala/com/cloudera/hue/livy/repl/Session.scala

@@ -16,11 +16,22 @@ object Session {
   case class Error() extends State
   case class ShuttingDown() extends State
   case class ShutDown() extends State
+
+  sealed trait Kind
+  case class Spark() extends Kind {
+    override def toString = "spark"
+  }
+
+  case class PySpark() extends Kind {
+    override def toString = "pyspark"
+  }
 }
 
 trait Session {
   import Session._
 
+  def kind: Kind
+
   def state: State
 
   def execute(code: String): Future[JValue]

apps/spark/java/livy-repl/src/main/scala/com/cloudera/hue/livy/repl/python/PythonSession.scala

@@ -195,6 +195,8 @@ private class PythonSession(process: Process, gatewayServer: GatewayServer) exte
 
   thread.start()
 
+  override def kind = Session.PySpark()
+
   override def state = _state
 
   override def history(): Seq[JValue] = _history

apps/spark/java/livy-repl/src/main/scala/com/cloudera/hue/livy/repl/scala/SparkSession.scala

@@ -22,6 +22,8 @@ private class SparkSession extends Session {
   private val interpreter = new Interpreter()
   interpreter.start()
 
+  override def kind: Session.Kind = Session.Spark()
+
   override def state: Session.State = interpreter.state match {
     case Interpreter.NotStarted() => Session.NotStarted()
     case Interpreter.Starting() => Session.Starting()

apps/spark/java/livy-repl/src/test/scala/com/cloudera/hue/livy/repl/WebAppSpec.scala

@@ -1,6 +1,6 @@
 package com.cloudera.hue.livy.repl
 
-import com.cloudera.hue.livy.repl.Session.State
+import com.cloudera.hue.livy.repl.Session.{Kind, State}
 import org.json4s.JsonAST.{JArray, JString}
 import org.json4s.JsonDSL._
 import org.json4s.jackson.JsonMethods._
@@ -17,6 +17,8 @@ class WebAppSpec extends ScalatraSuite with FunSpecLike with BeforeAndAfter {
     var _state: State = Session.Idle()
     var _history = List[JValue]()
 
+    override def kind: Kind = Session.Spark()
+
     override def state = _state
 
     override def execute(code: String): Future[JValue] = {

apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/SessionFactory.scala

@@ -9,7 +9,7 @@ import com.cloudera.hue.livy.yarn.Client
 import scala.concurrent.{ExecutionContext, Future}
 
 trait SessionFactory {
-  def createSession(lang: String): Future[Session]
+  def createSession(kind: Session.Kind): Future[Session]
 
   def close(): Unit = {}
 }
@@ -18,10 +18,10 @@ class ThreadSessionFactory(livyConf: LivyConf) extends SessionFactory {
 
   implicit def executor: ExecutionContext = ExecutionContext.global
 
-  override def createSession(lang: String): Future[Session] = {
+  override def createSession(kind: Session.Kind): Future[Session] = {
     Future {
       val id = UUID.randomUUID().toString
-      ThreadSession.create(id, lang)
+      ThreadSession.create(id, kind)
     }
   }
 }
@@ -30,10 +30,10 @@ class ProcessSessionFactory(livyConf: LivyConf) extends SessionFactory {
 
   implicit def executor: ExecutionContext = ExecutionContext.global
 
-  override def createSession(lang: String): Future[Session] = {
+  override def createSession(kind: Session.Kind): Future[Session] = {
     Future {
       val id = UUID.randomUUID().toString
-      ProcessSession.create(livyConf, id, lang)
+      ProcessSession.create(livyConf, id, kind)
     }
   }
 }
@@ -42,9 +42,9 @@ class YarnSessionFactory(livyConf: LivyConf) extends SessionFactory {
 
   val client = new Client(livyConf)
 
-  override def createSession(lang: String): Future[Session] = {
+  override def createSession(kind: Session.Kind): Future[Session] = {
     val id = UUID.randomUUID().toString
-    YarnSession.create(client, id, lang)
+    YarnSession.create(client, id, kind)
   }
 
   override def close(): Unit = {

apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/SessionManager.scala

@@ -32,8 +32,8 @@ class SessionManager(factory: SessionFactory) extends Logging {
     sessions.keys
   }
 
-  def createSession(lang: String): Future[Session] = {
-    val session = factory.createSession(lang)
+  def createSession(kind: Session.Kind): Future[Session] = {
+    val session = factory.createSession(kind)
 
     session.map({ case(session: Session) =>
       info("created session %s" format session.id)

apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/WebApp.scala

@@ -54,10 +54,8 @@ class WebApp(sessionManager: SessionManager)
     val createSessionRequest = parsedBody.extract[CreateSessionRequest]
 
     val sessionFuture = createSessionRequest.lang match {
-      case "scala" => sessionManager.createSession(createSessionRequest.lang)
-      case "spark" => sessionManager.createSession(createSessionRequest.lang)
-      case "pyspark" => sessionManager.createSession(createSessionRequest.lang)
-      case "python" => sessionManager.createSession(createSessionRequest.lang)
+      case "spark" | "scala" => sessionManager.createSession(Session.Spark())
+      case "pyspark" | "python" => sessionManager.createSession(Session.PySpark())
       case lang => halt(400, "unsupported language: " + lang)
     }
 
@@ -166,6 +164,7 @@ class WebApp(sessionManager: SessionManager)
   private def formatSession(session: Session) = {
     Map(
       "id" -> session.id,
+      "kind" -> session.kind.toString,
       "state" -> session.state.getClass.getSimpleName.toLowerCase
     )
   }

apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/sessions/ProcessSession.scala

@@ -18,17 +18,16 @@ object ProcessSession extends Logging {
   val CONF_LIVY_REPL_CALLBACK_URL = "livy.repl.callback-url"
   val CONF_LIVY_REPL_DRIVER_CLASS_PATH = "livy.repl.driverClassPath"
 
-  def create(livyConf: LivyConf, id: String, lang: String): Session = {
-    val process = startProcess(livyConf, id, lang)
-    new ProcessSession(id, process)
+  def create(livyConf: LivyConf, id: String, kind: Session.Kind): Session = {
+    val process = startProcess(livyConf, id, kind)
+    new ProcessSession(id, kind, process)
   }
 
   // Loop until we've started a process with a valid port.
-  private def startProcess(livyConf: LivyConf, id: String, lang: String): Process = {
+  private def startProcess(livyConf: LivyConf, id: String, kind: Session.Kind): Process = {
     val args = ArrayBuffer(
       "spark-submit",
-      "--class",
-      "com.cloudera.hue.livy.repl.Main"
+      "--class", "com.cloudera.hue.livy.repl.Main"
     )
 
     sys.env.get("LIVY_REPL_JAVA_OPTS").foreach { case javaOpts =>
@@ -42,7 +41,7 @@ object ProcessSession extends Logging {
     }
 
     args += livyJar(livyConf)
-    args += lang
+    args += kind.toString
 
     val pb = new ProcessBuilder(args)
 
@@ -65,7 +64,7 @@ object ProcessSession extends Logging {
   }
 }
 
-private class ProcessSession(id: String, process: Process) extends WebSession(id) {
+private class ProcessSession(id: String, kind: Session.Kind, process: Process) extends WebSession(id, kind) {
 
   val stdoutThread = new Thread {
     override def run() = {

apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/sessions/Session.scala

@@ -12,12 +12,40 @@ import scala.concurrent.duration.Duration
 
 object Session {
   sealed trait State
-  case class NotStarted() extends State
-  case class Starting() extends State
-  case class Idle() extends State
-  case class Busy() extends State
-  case class Error() extends State
-  case class Dead() extends State
+
+  case class NotStarted() extends State {
+    override def toString = "not_started"
+  }
+
+  case class Starting() extends State {
+    override def toString = "starting"
+  }
+
+  case class Idle() extends State {
+    override def toString = "idle"
+  }
+
+  case class Busy() extends State {
+    override def toString = "busy"
+  }
+
+  case class Error() extends State {
+    override def toString = "error"
+  }
+
+  case class Dead() extends State {
+    override def toString = "dead"
+  }
+
+  sealed trait Kind
+
+  case class Spark() extends Kind {
+    override def toString = "spark"
+  }
+
+  case class PySpark() extends Kind {
+    override def toString = "pyspark"
+  }
 
   class SessionFailedToStart(msg: String) extends Exception(msg)
 
@@ -29,6 +57,8 @@ trait Session {
 
   def id: String
 
+  def kind: Kind
+
   def lastActivity: Long
 
   def state: State

apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/sessions/ThreadSession.scala

@@ -16,18 +16,18 @@ object ThreadSession {
   val LIVY_HOME = System.getenv("LIVY_HOME")
   val LIVY_REPL = LIVY_HOME + "/bin/livy-repl"
 
-  def create(id: String, lang: String): Session = {
-    val session = lang match {
-      case "python" | "pyspark" =>
-        PythonSession.createPySpark()
-      case "scala" | "spark" =>
+  def create(id: String, kind: Session.Kind): Session = {
+    val session = kind match {
+      case Session.Spark() =>
         SparkSession.create()
+      case Session.PySpark() =>
+        PythonSession.createPySpark()
     }
-    new ThreadSession(id, session)
+    new ThreadSession(id, kind, session)
   }
 }
 
-private class ThreadSession(val id: String, session: com.cloudera.hue.livy.repl.Session) extends Session {
+private class ThreadSession(val id: String, val kind: Session.Kind, session: com.cloudera.hue.livy.repl.Session) extends Session {
 
   protected implicit def executor: ExecutionContextExecutor = ExecutionContext.global
 

apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/sessions/WebSession.scala

@@ -9,14 +9,14 @@ import com.cloudera.hue.livy.server.Statement
 import com.cloudera.hue.livy.server.sessions.Session._
 import dispatch._
 import org.json4s.jackson.Serialization.write
-import org.json4s.{DefaultFormats, Formats}
+import org.json4s.{JValue, DefaultFormats, Formats}
 
 import scala.annotation.tailrec
 import scala.collection.mutable.ArrayBuffer
 import scala.concurrent.duration.Duration
 import scala.concurrent.{Future, _}
 
-class WebSession(val id: String) extends Session with Logging {
+class WebSession(val id: String, val kind: Kind) extends Session with Logging {
 
   protected implicit def executor: ExecutionContextExecutor = ExecutionContext.global
   protected implicit def jsonFormats: Formats = DefaultFormats

apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/sessions/YarnSession.scala

@@ -10,15 +10,15 @@ import scala.concurrent.duration._
 object YarnSession {
   protected implicit def executor: ExecutionContextExecutor = ExecutionContext.global
 
-  def create(client: Client, id: String, lang: String): Future[Session] = {
+  def create(client: Client, id: String, kind: Session.Kind): Future[Session] = {
     val callbackUrl = System.getProperty("livy.server.callback-url")
-    val job = client.submitApplication(id, lang, callbackUrl)
+    val job = client.submitApplication(id, kind.toString, callbackUrl)
 
-    Future.successful(new YarnSession(id, job))
+    Future.successful(new YarnSession(id, kind, job))
   }
 }
 
-private class YarnSession(id: String, job: Future[Job]) extends WebSession(id) {
+private class YarnSession(id: String, kind: Session.Kind, job: Future[Job]) extends WebSession(id, kind) {
   job.onFailure { case _ =>
     _state = Session.Error()
   }

apps/spark/java/livy-server/src/test/scala/com/cloudera/hue/livy/server/ProcessSessionSpec.scala

@@ -10,5 +10,5 @@ class ProcessSessionSpec extends BaseSessionSpec with FunSpecLike with Matchers
   val livyConf = new LivyConf()
   livyConf.set("livy.repl.driverClassPath", sys.props("java.class.path"))
 
-  def createSession() = ProcessSession.create(livyConf, "0", "spark")
+  def createSession() = ProcessSession.create(livyConf, "0", Session.Spark())
 }

apps/spark/java/livy-server/src/test/scala/com/cloudera/hue/livy/server/ThreadSessionSpec.scala

@@ -1,9 +1,9 @@
 package com.cloudera.hue.livy.server
 
-import com.cloudera.hue.livy.server.sessions.ThreadSession
+import com.cloudera.hue.livy.server.sessions.{Session, ThreadSession}
 import org.scalatest.{BeforeAndAfter, FunSpecLike, Matchers}
 
 class ThreadSessionSpec extends BaseSessionSpec with FunSpecLike with Matchers with BeforeAndAfter {
 
-  def createSession() = ThreadSession.create("0", "spark")
+  def createSession() = ThreadSession.create("0", Session.Spark())
 }

apps/spark/java/livy-yarn/src/main/scala/com/cloudera/hue/livy/yarn/Client.scala

@@ -39,7 +39,7 @@ class Client(livyConf: LivyConf) extends Logging {
   yarnClient.init(yarnConf)
   yarnClient.start()
 
-  def submitApplication(id: String, lang: String, callbackUrl: String): Future[Job] = {
+  def submitApplication(id: String, kind: String, callbackUrl: String): Future[Job] = {
     val url = f"$callbackUrl/sessions/$id/callback"
 
     val builder: ProcessBuilder = new ProcessBuilder(
@@ -48,7 +48,7 @@ class Client(livyConf: LivyConf) extends Logging {
       "--class", "com.cloudera.hue.livy.repl.Main",
       "--driver-java-options", f"-Dlivy.repl.callback-url=$url -Dlivy.repl.port=0",
       livyJar(livyConf),
-      lang
+      kind
     )
 
     builder.redirectOutput(Redirect.PIPE)