浏览代码

[livy] Support paging in livy-repl

Erick Tryzelaar 10 年之前
父节点
当前提交
99db83f

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

@@ -30,11 +30,9 @@ trait Session {
 
   def state: State
 
-  def execute(code: String): Future[JValue]
+  def execute(code: String): Statement
 
-  def history(): Seq[JValue]
-
-  def history(id: Int): Option[JValue]
+  def history: IndexedSeq[Statement]
 
   def close(): Unit
 
@@ -44,3 +42,5 @@ trait Session {
     Utils.waitUntil({ () => state != oldState }, atMost)
   }
 }
+
+case class Statement(id: Int, result: Future[JValue])

+ 81 - 20
apps/spark/java/livy-repl/src/main/scala/com/cloudera/hue/livy/repl/WebApp.scala

@@ -18,22 +18,25 @@
 
 package com.cloudera.hue.livy.repl
 
+import java.util.concurrent.TimeUnit
+
 import com.cloudera.hue.livy.Logging
 import com.cloudera.hue.livy.msgs.ExecuteRequest
 import com.cloudera.hue.livy.sessions._
 import com.fasterxml.jackson.core.JsonParseException
-import org.json4s.{DefaultFormats, MappingException}
+import org.json4s.{DefaultFormats, JsonDSL, MappingException, _}
 import org.scalatra._
 import org.scalatra.json.JacksonJsonSupport
 
-import _root_.scala.concurrent.{Future, ExecutionContext}
+import _root_.scala.concurrent.duration.Duration
+import _root_.scala.concurrent.{Await, ExecutionContext, Future, TimeoutException}
 
 object WebApp extends Logging
 
 class WebApp(session: Session) extends ScalatraServlet with FutureSupport with JacksonJsonSupport {
 
   override protected implicit def executor: ExecutionContext = ExecutionContext.global
-  override protected implicit val jsonFormats = DefaultFormats
+  override protected implicit def jsonFormats: Formats = DefaultFormats ++ Serializers.Formats
 
   before() {
     contentType = formats("json")
@@ -45,35 +48,26 @@ class WebApp(session: Session) extends ScalatraServlet with FutureSupport with J
   }
 
   get("/") {
-    val state = session.state match {
-      case NotStarted() => "not_started"
-      case Starting() => "starting"
-      case Idle() => "idle"
-      case Busy() => "busy"
-      case Running() => "running"
-      case Error() => "error"
-      case ShuttingDown() => "shutting_down"
-      case Dead() => "dead"
-      case Success() => "success"
-    }
-    Map("state" -> state)
+    Serializers.serializeSession(session)
   }
 
   post("/execute") {
     val req = parsedBody.extract[ExecuteRequest]
-    val rep = session.execute(req.code)
-    new AsyncResult { val is = rep }
+    Serializers.serializeStatement(session.execute(req.code))
   }
 
   get("/history") {
-    session.history()
+    val from = params.get("from").map(_.toInt)
+    val size = params.get("size").map(_.toInt)
+
+    Serializers.serializeHistory(session.history, from, size)
   }
 
   get("/history/:statementId") {
     val statementId = params("statementId").toInt
 
-    session.history(statementId) match {
-      case Some(statement) => statement
+    session.history.lift(statementId) match {
+      case Some(statement) => Serializers.serializeStatement(statement)
       case None => NotFound("Statement not found")
     }
   }
@@ -95,3 +89,70 @@ class WebApp(session: Session) extends ScalatraServlet with FutureSupport with J
   }
 }
 
+private object Serializers {
+  import JsonDSL._
+
+  def Formats: List[CustomSerializer[_]] = List(StatementSerializer)
+
+  def serializeSession(session: Session): JValue = {
+    val state = session.state match {
+      case NotStarted() => "not_started"
+      case Starting() => "starting"
+      case Idle() => "idle"
+      case Busy() => "busy"
+      case Running() => "running"
+      case Error() => "error"
+      case ShuttingDown() => "shutting_down"
+      case Dead() => "dead"
+      case Success() => "success"
+    }
+    Map("state" -> state)
+  }
+
+  def serializeHistory(history: IndexedSeq[Statement],
+                       fromOpt: Option[Int],
+                       sizeOpt: Option[Int]): JValue = {
+    val size = sizeOpt.getOrElse(100)
+    var from = fromOpt.getOrElse(-1)
+    if (from < 0) {
+      from = math.max(0, history.length - size)
+    }
+    val until = from + size
+
+    ("from", from) ~
+      ("total", history.length) ~
+      ("statements", history.view(from, until)
+        .map(serializeStatement(_, Duration(0, TimeUnit.SECONDS))))
+  }
+
+  def serializeStatement(statement: Statement,
+                         duration: Duration = Duration(10, TimeUnit.SECONDS)): JValue = {
+    val result = try {
+      Await.result(statement.result, duration)
+    } catch {
+      case _: TimeoutException => null
+    }
+
+    ("id", statement.id) ~ ("result", result)
+  }
+
+  case object SessionSerializer extends CustomSerializer[Session](
+    implicit formats => ( {
+      // We don't support deserialization.
+      PartialFunction.empty
+    }, {
+      case session: Session => serializeSession(session)
+    }
+      )
+  )
+
+  case object StatementSerializer extends CustomSerializer[Statement](
+    implicit formats => ( {
+      // We don't support deserialization.
+      PartialFunction.empty
+    }, {
+      case statement: Statement => serializeStatement(statement)
+    }
+      )
+  )
+}

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

@@ -19,11 +19,10 @@
 package com.cloudera.hue.livy.repl.python
 
 import com.cloudera.hue.livy.Logging
-import com.cloudera.hue.livy.repl.{Interpreter, Session}
+import com.cloudera.hue.livy.repl.{Interpreter, Session, Statement}
 import com.cloudera.hue.livy.sessions._
 import org.json4s.JValue
 
-import scala.collection.mutable.ArrayBuffer
 import scala.concurrent._
 
 object PythonSession {
@@ -35,30 +34,19 @@ object PythonSession {
 private class PythonSession(interpreter: Interpreter) extends Session with Logging {
   private implicit def executor: ExecutionContext = ExecutionContext.global
 
-  private var _history = ArrayBuffer[JValue]()
+  private var _history = IndexedSeq[Statement]()
 
   override def kind = PySpark()
 
   override def state = interpreter.state
 
-  override def history(): Seq[JValue] = _history
+  override def history: IndexedSeq[Statement] = _history
 
-  override def history(id: Int): Option[JValue] = {
-    if (id < _history.length) {
-      Some(_history(id))
-    } else {
-      None
-    }
-  }
-
-  override def execute(code: String): Future[JValue] = {
-    val future = interpreter.execute(code)
-    future.onSuccess { case result =>
-      synchronized {
-        _history += result
-      }
-    }
-    future
+  override def execute(code: String): Statement = synchronized {
+    val result = interpreter.execute(code)
+    val statement = Statement(_history.length, result)
+    _history :+= statement
+    statement
   }
 
   override def close(): Unit = interpreter.close()

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

@@ -18,7 +18,7 @@
 
 package com.cloudera.hue.livy.repl.scala
 
-import com.cloudera.hue.livy.repl.Session
+import com.cloudera.hue.livy.repl.{Statement, Session}
 import com.cloudera.hue.livy.repl.scala.interpreter._
 import com.cloudera.hue.livy.sessions._
 import org.json4s.jackson.JsonMethods._
@@ -37,7 +37,7 @@ private class SparkSession extends Session {
 
   implicit val formats = DefaultFormats
 
-  private var _history = new mutable.ArrayBuffer[JValue]
+  private var _history = IndexedSeq[Statement]()
   private val interpreter = new Interpreter()
   interpreter.start()
 
@@ -52,18 +52,10 @@ private class SparkSession extends Session {
     case Interpreter.ShutDown() => Dead()
   }
 
-  override def history(): Seq[JValue] = _history
+  override def history: IndexedSeq[Statement] = _history
 
-  override def history(id: Int): Option[JValue] = synchronized {
-    if (id < _history.length) {
-      Some(_history(id))
-    } else {
-      None
-    }
-  }
-
-  override def execute(code: String): Future[JValue] = {
-    Future {
+  override def execute(code: String): Statement = synchronized {
+    val result = Future {
       val content = interpreter.execute(code) match {
         case ExecuteComplete(executeCount, output) =>
           Map(
@@ -89,12 +81,12 @@ private class SparkSession extends Session {
           )
       }
 
-      val jsonContent = parse(write(content))
-
-      _history += jsonContent
-
-      jsonContent
+      parse(write(content))
     }
+
+    val statement = Statement(_history.length, result)
+    _history :+= statement
+    statement
   }
 
   override def close(): Unit = {

+ 42 - 13
apps/spark/java/livy-repl/src/test/scala/com/cloudera/hue/livy/repl/PythonSessionSpec.scala

@@ -31,7 +31,10 @@ class PythonSessionSpec extends BaseSessionSpec {
 
   describe("A python session") {
     it("should execute `1 + 2` == 3") {
-      val result = Await.result(session.execute("1 + 2"), Duration.Inf)
+      val statement = session.execute("1 + 2")
+      statement.id should equal (0)
+
+      val result = Await.result(statement.result, Duration.Inf)
       val expectedResult = Extraction.decompose(Map(
         "status" -> "ok",
         "execution_count" -> 0,
@@ -44,7 +47,10 @@ class PythonSessionSpec extends BaseSessionSpec {
     }
 
     it("should execute `x = 1`, then `y = 2`, then `x + y`") {
-      var result = Await.result(session.execute("x = 1"), Duration.Inf)
+      var statement = session.execute("x = 1")
+      statement.id should equal (0)
+
+      var result = Await.result(statement.result, Duration.Inf)
       var expectedResult = Extraction.decompose(Map(
         "status" -> "ok",
         "execution_count" -> 0,
@@ -55,7 +61,10 @@ class PythonSessionSpec extends BaseSessionSpec {
 
       result should equal (expectedResult)
 
-      result = Await.result(session.execute("y = 2"), Duration.Inf)
+      statement = session.execute("y = 2")
+      statement.id should equal (1)
+
+      result = Await.result(statement.result, Duration.Inf)
       expectedResult = Extraction.decompose(Map(
         "status" -> "ok",
         "execution_count" -> 1,
@@ -66,7 +75,10 @@ class PythonSessionSpec extends BaseSessionSpec {
 
       result should equal (expectedResult)
 
-      result = Await.result(session.execute("x + y"), Duration.Inf)
+      statement = session.execute("x + y")
+      statement.id should equal (2)
+
+      result = Await.result(statement.result, Duration.Inf)
       expectedResult = Extraction.decompose(Map(
         "status" -> "ok",
         "execution_count" -> 2,
@@ -79,7 +91,10 @@ class PythonSessionSpec extends BaseSessionSpec {
     }
 
     it("should do table magic") {
-      val result = Await.result(session.execute("x = [[1, 'a'], [3, 'b']]\n%table x"), Duration.Inf)
+      val statement = session.execute("x = [[1, 'a'], [3, 'b']]\n%table x")
+      statement.id should equal (0)
+
+      val result = Await.result(statement.result, Duration.Inf)
       val expectedResult = Extraction.decompose(Map(
         "status" -> "ok",
         "execution_count" -> 1,
@@ -97,7 +112,10 @@ class PythonSessionSpec extends BaseSessionSpec {
     }
 
     it("should capture stdout") {
-      val result = Await.result(session.execute("""print 'Hello World'"""), Duration.Inf)
+      val statement = session.execute("""print 'Hello World'""")
+      statement.id should equal (0)
+
+      val result = Await.result(statement.result, Duration.Inf)
       val expectedResult = Extraction.decompose(Map(
         "status" -> "ok",
         "execution_count" -> 0,
@@ -110,7 +128,10 @@ class PythonSessionSpec extends BaseSessionSpec {
     }
 
     it("should report an error if accessing an unknown variable") {
-      val result = Await.result(session.execute("""x"""), Duration.Inf)
+      val statement = session.execute("""x""")
+      statement.id should equal (0)
+
+      val result = Await.result(statement.result, Duration.Inf)
       val expectedResult = Extraction.decompose(Map(
         "status" -> "error",
         "execution_count" -> 0,
@@ -126,11 +147,14 @@ class PythonSessionSpec extends BaseSessionSpec {
     }
 
     it("should report an error if exception is thrown") {
-      val result = Await.result(session.execute(
+      val statement = session.execute(
         """def foo():
           |    raise Exception()
           |foo()
-          |""".stripMargin), Duration.Inf)
+          |""".stripMargin)
+      statement.id should equal (0)
+
+      val result = Await.result(statement.result, Duration.Inf)
       val expectedResult = Extraction.decompose(Map(
         "status" -> "error",
         "execution_count" -> 0,
@@ -146,7 +170,10 @@ class PythonSessionSpec extends BaseSessionSpec {
     }
 
     it("should access the spark context") {
-      val result = Await.result(session.execute("""sc"""), Duration.Inf)
+      val statement = session.execute("""sc""")
+      statement.id should equal (0)
+
+      val result = Await.result(statement.result, Duration.Inf)
       val resultMap = result.extract[Map[String, JValue]]
 
       // Manually extract the values since the line numbers in the exception could change.
@@ -158,10 +185,12 @@ class PythonSessionSpec extends BaseSessionSpec {
     }
 
     it("should execute spark commands") {
-      val result = Await.result(session.execute(
-        """
+      val statement = session.execute("""
           |sc.parallelize(xrange(0, 2)).map(lambda i: i + 1).collect()
-          |""".stripMargin), Duration.Inf)
+          |""".stripMargin)
+      statement.id should equal (0)
+
+      val result = Await.result(statement.result, Duration.Inf)
 
       val expectedResult = Extraction.decompose(Map(
         "status" -> "ok",

+ 37 - 10
apps/spark/java/livy-repl/src/test/scala/com/cloudera/hue/livy/repl/SparkSessionSpec.scala

@@ -31,7 +31,10 @@ class SparkSessionSpec extends BaseSessionSpec {
 
   describe("A spark session") {
     it("should execute `1 + 2` == 3") {
-      val result = Await.result(session.execute("1 + 2"), Duration.Inf)
+      val statement = session.execute("1 + 2")
+      statement.id should equal (0)
+
+      val result = Await.result(statement.result, Duration.Inf)
       val expectedResult = Extraction.decompose(Map(
         "status" -> "ok",
         "execution_count" -> 0,
@@ -44,7 +47,10 @@ class SparkSessionSpec extends BaseSessionSpec {
     }
 
     it("should execute `x = 1`, then `y = 2`, then `x + y`") {
-      var result = Await.result(session.execute("val x = 1"), Duration.Inf)
+      var statement = session.execute("val x = 1")
+      statement.id should equal (0)
+
+      var result = Await.result(statement.result, Duration.Inf)
       var expectedResult = Extraction.decompose(Map(
         "status" -> "ok",
         "execution_count" -> 0,
@@ -55,7 +61,10 @@ class SparkSessionSpec extends BaseSessionSpec {
 
       result should equal (expectedResult)
 
-      result = Await.result(session.execute("val y = 2"), Duration.Inf)
+      statement = session.execute("val y = 2")
+      statement.id should equal (1)
+
+      result = Await.result(statement.result, Duration.Inf)
       expectedResult = Extraction.decompose(Map(
         "status" -> "ok",
         "execution_count" -> 1,
@@ -66,7 +75,10 @@ class SparkSessionSpec extends BaseSessionSpec {
 
       result should equal (expectedResult)
 
-      result = Await.result(session.execute("x + y"), Duration.Inf)
+      statement = session.execute("x + y")
+      statement.id should equal (2)
+
+      result = Await.result(statement.result, Duration.Inf)
       expectedResult = Extraction.decompose(Map(
         "status" -> "ok",
         "execution_count" -> 2,
@@ -79,7 +91,10 @@ class SparkSessionSpec extends BaseSessionSpec {
     }
 
     it("should capture stdout") {
-      val result = Await.result(session.execute("""println("Hello World")"""), Duration.Inf)
+      val statement = session.execute("""println("Hello World")""")
+      statement.id should equal (0)
+
+      val result = Await.result(statement.result, Duration.Inf)
       val expectedResult = Extraction.decompose(Map(
         "status" -> "ok",
         "execution_count" -> 0,
@@ -92,7 +107,10 @@ class SparkSessionSpec extends BaseSessionSpec {
     }
 
     it("should report an error if accessing an unknown variable") {
-      val result = Await.result(session.execute("""x"""), Duration.Inf)
+      val statement = session.execute("""x""")
+      statement.id should equal (0)
+
+      val result = Await.result(statement.result, Duration.Inf)
       val expectedResult = Extraction.decompose(Map(
         "status" -> "error",
         "execution_count" -> 0,
@@ -107,7 +125,10 @@ class SparkSessionSpec extends BaseSessionSpec {
     }
 
     it("should report an error if exception is thrown") {
-      val result = Await.result(session.execute("""throw new Exception()"""), Duration.Inf)
+      val statement = session.execute("""throw new Exception()""")
+      statement.id should equal (0)
+
+      val result = Await.result(statement.result, Duration.Inf)
       val resultMap = result.extract[Map[String, JValue]]
 
       // Manually extract the values since the line numbers in the exception could change.
@@ -119,7 +140,10 @@ class SparkSessionSpec extends BaseSessionSpec {
     }
 
     it("should access the spark context") {
-      val result = Await.result(session.execute("""sc"""), Duration.Inf)
+      val statement = session.execute("""sc""")
+      statement.id should equal (0)
+
+      val result = Await.result(statement.result, Duration.Inf)
       val resultMap = result.extract[Map[String, JValue]]
 
       // Manually extract the values since the line numbers in the exception could change.
@@ -131,10 +155,13 @@ class SparkSessionSpec extends BaseSessionSpec {
     }
 
     it("should execute spark commands") {
-      val result = Await.result(session.execute(
+      val statement = session.execute(
         """
           |sc.parallelize(0 to 1).map{i => i+1}.collect
-          |""".stripMargin), Duration.Inf)
+          |""".stripMargin)
+      statement.id should equal (0)
+
+      val result = Await.result(statement.result, Duration.Inf)
 
       val expectedResult = Extraction.decompose(Map(
         "status" -> "ok",

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

@@ -22,9 +22,10 @@ import com.cloudera.hue.livy.sessions._
 import org.json4s.JsonAST.{JArray, JString}
 import org.json4s.JsonDSL._
 import org.json4s.jackson.JsonMethods._
-import org.json4s.{DefaultFormats, Extraction, JValue}
-import org.scalatest.{FunSpecLike, BeforeAndAfter}
+import org.json4s.{DefaultFormats, Extraction}
+import org.scalatest.{BeforeAndAfter, FunSpecLike}
 import org.scalatra.test.scalatest.ScalatraSuite
+
 import _root_.scala.concurrent.Future
 
 class WebAppSpec extends ScalatraSuite with FunSpecLike with BeforeAndAfter {
@@ -33,24 +34,24 @@ class WebAppSpec extends ScalatraSuite with FunSpecLike with BeforeAndAfter {
 
   class MockSession extends Session {
     var _state: State = Idle()
-    var _history = List[JValue]()
+    var _history = IndexedSeq[Statement]()
 
     override def kind: Kind = Spark()
 
     override def state = _state
 
-    override def execute(code: String): Future[JValue] = {
+    override def execute(code: String): Statement = {
       val rep = render(Map("hi" -> "there"))
-      Future.successful(rep)
+      val statement = Statement(0, Future.successful(rep))
+      _history :+= statement
+      statement
     }
 
     override def close(): Unit = {
       _state = Dead()
     }
 
-    override def history(): Seq[JValue] = _history
-
-    override def history(id: Int): Option[JValue] = _history.lift(id)
+    override def history: IndexedSeq[Statement] = _history
   }
 
   val session = new MockSession
@@ -81,24 +82,32 @@ class WebAppSpec extends ScalatraSuite with FunSpecLike with BeforeAndAfter {
       get("/history") {
         status should equal (200)
         header("Content-Type") should include("application/json")
-        parse(body) should equal (JArray(List()))
+        parse(body) should equal (
+          ("from", 0) ~
+            ("total", 0) ~
+            ("statements", JArray(List())))
       }
     }
 
     it("GET /history with history should return something") {
-      val history = Extraction.decompose(Map("data" -> Map("text/plain" -> "1")))
-      session._history = List(history)
+      val history = Extraction.decompose(Map(
+          "data" -> Map("text/plain" -> "1")
+      ))
+      session._history = IndexedSeq(Statement(0, Future.successful(history)))
 
       get("/history") {
         status should equal (200)
         header("Content-Type") should include("application/json")
-        parse(body) should equal (JArray(List(history)))
+        parse(body) should equal (
+          ("from", 0) ~
+            ("total", 1) ~
+            ("statements", JArray(List( ("id", 0) ~ ("result", history)))))
       }
     }
   }
 
   after {
     session._state = Idle()
-    session._history = List()
+    session._history = IndexedSeq()
   }
 }

+ 3 - 3
apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/interactive/InteractiveWebSession.scala

@@ -26,7 +26,7 @@ import com.cloudera.hue.livy.msgs.ExecuteRequest
 import com.cloudera.hue.livy.sessions._
 import dispatch._
 import org.json4s.jackson.Serialization.write
-import org.json4s.{DefaultFormats, Formats}
+import org.json4s.{DefaultFormats, Formats, JValue}
 
 import scala.concurrent.duration.Duration
 import scala.concurrent.{Future, _}
@@ -74,10 +74,10 @@ abstract class InteractiveWebSession(val id: Int, createInteractiveRequest: Crea
       var req = (svc / "execute").setContentType("application/json", "UTF-8")
       req = req << write(content)
 
-      val future = Http(req OK as.json4s.Json).map { case (resp) =>
+      val future = Http(req OK as.json4s.Json).map { case resp: JValue =>
         synchronized {
           transition(Idle())
-          resp
+          resp \ "result"
         }
       }