Browse Source

HUE-2591 [livy] Add support for paging through sessions, statements, and statement output

To support paging, add the query parameters "from" and "size" to the
arguments. By default it fetches the first 100 items.
Erick Tryzelaar 10 years ago
parent
commit
8370286b7f

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

@@ -51,11 +51,7 @@ trait Session {
 
   def executeStatement(content: ExecuteRequest): Statement
 
-  def statement(statementId: Int): Option[Statement]
-
-  def statements(): Seq[Statement]
-
-  def statements(fromIndex: Integer, toIndex: Integer): Seq[Statement]
+  def statements: IndexedSeq[Statement]
 
   def interrupt(): Future[Unit]
 

+ 35 - 19
apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/sessions/SessionServlet.scala

@@ -135,11 +135,15 @@ class SessionServlet(sessionManager: SessionManager)
     val sessionId = params("sessionId").toInt
 
     sessionManager.get(sessionId) match {
+      case None => NotFound("Session not found")
       case Some(session: Session) =>
+        val from = params.get("from").map(_.toInt).getOrElse(0)
+        val size = params.get("size").map(_.toInt).getOrElse(session.statements.length)
+
         Map(
-          "statements" -> session.statements()
+          "total_statements" -> session.statements.length,
+          "statements" -> session.statements.view(from, from + size)
         )
-      case None => NotFound("Session not found")
     }
   }
 
@@ -147,13 +151,17 @@ class SessionServlet(sessionManager: SessionManager)
     val sessionId = params("sessionId").toInt
     val statementId = params("statementId").toInt
 
+    val from = params.get("from").map(_.toInt)
+    val size = params.get("size").map(_.toInt)
+
     sessionManager.get(sessionId) match {
+      case None => NotFound("Session not found")
       case Some(session) =>
-        session.statement(statementId) match {
-          case Some(statement) => statement
+        session.statements.lift(statementId) match {
           case None => NotFound("Statement not found")
+          case Some(statement) =>
+            Serializers.serializeStatement(statement, from, size)
         }
-      case None => NotFound("Session not found")
     }
   }
 
@@ -201,15 +209,32 @@ private object Serializers {
 
   private def serializeStatementState(state: Statement.State) = JString(state.toString)
 
+  def serializeSession(session: Session): JValue = {
+    ("id", session.id) ~
+      ("state", serializeSessionState(session.state)) ~
+      ("kind", serializeSessionKind(session.kind)) ~
+      ("proxyUser", session.proxyUser)
+  }
+
+  def serializeStatement(statement: Statement, from: Option[Int], size: Option[Int]): JValue = {
+    // Take a couple milliseconds to see if the statement has finished.
+    val output = try {
+      Await.result(statement.output(), Duration(100, TimeUnit.MILLISECONDS))
+    } catch {
+      case _: TimeoutException => null
+    }
+
+    ("id" -> statement.id) ~
+      ("state" -> serializeStatementState(statement.state)) ~
+      ("output" -> output)
+  }
+
   case object SessionSerializer extends CustomSerializer[Session](implicit formats => ( {
     // We don't support deserialization.
     PartialFunction.empty
   }, {
     case session: Session =>
-      ("id", session.id) ~
-      ("state", serializeSessionState(session.state)) ~
-      ("kind", serializeSessionKind(session.kind)) ~
-      ("proxyUser", session.proxyUser)
+      serializeSession(session)
   }
     )
   )
@@ -237,16 +262,7 @@ private object Serializers {
     PartialFunction.empty
   }, {
     case statement: Statement =>
-      // Take a couple milliseconds to see if the statement has finished.
-      val output = try {
-        Await.result(statement.output, Duration(100, TimeUnit.MILLISECONDS))
-      } catch {
-        case _: TimeoutException => null
-      }
-
-      ("id" -> statement.id) ~
-        ("state" -> serializeStatementState(statement.state)) ~
-        ("output" -> output)
+      serializeStatement(statement, None, None)
   }))
 
   case object StatementStateSerializer extends CustomSerializer[Statement.State](implicit formats => ( {

+ 29 - 2
apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/sessions/Statement.scala

@@ -20,6 +20,7 @@ package com.cloudera.hue.livy.server.sessions
 
 import com.cloudera.hue.livy.msgs.ExecuteRequest
 import org.json4s.JValue
+import org.json4s.JsonAST.{JArray, JObject, JField, JString}
 
 import scala.concurrent.{ExecutionContext, ExecutionContextExecutor, Future}
 import scala.util.{Failure, Success}
@@ -40,7 +41,7 @@ object Statement {
   }
 }
 
-class Statement(val id: Int, val request: ExecuteRequest, val output: Future[JValue]) {
+class Statement(val id: Int, val request: ExecuteRequest, _output: Future[JValue]) {
   import Statement._
 
   protected implicit def executor: ExecutionContextExecutor = ExecutionContext.global
@@ -49,7 +50,33 @@ class Statement(val id: Int, val request: ExecuteRequest, val output: Future[JVa
 
   def state = _state
 
-  output.onComplete {
+  def output(from: Option[Int] = None, size: Option[Int] = None): Future[JValue] = {
+    _output.map { case output =>
+      if (from.isEmpty && size.isEmpty) {
+        output
+      } else {
+        val from_ = from.getOrElse(0)
+        val size_ = size.getOrElse(100)
+        val until = from_ + size_
+
+        output \ "data" match {
+          case JObject(JField("text/plain", JString(text)) :: Nil) =>
+            val lines = text.split('\n').slice(from_, until)
+            output.replace(
+              "data" :: "text/plain" :: Nil,
+              JString(lines.mkString("\n")))
+          case JObject(JField("application/json", JArray(items)) :: Nil) =>
+            output.replace(
+              "data" :: "application/json" :: Nil,
+              JArray(items.slice(from_, until)))
+          case _ =>
+            output
+        }
+      }
+    }
+  }
+
+  _output.onComplete {
     case Success(_) => _state = Available()
     case Failure(_) => _state = Error()
   }

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

@@ -71,11 +71,7 @@ private class ThreadSession(val id: Int,
     statement
   }
 
-  override def statement(statementId: Int): Option[Statement] = statements_.lift(statementId)
-
-  override def statements(): Seq[Statement] = statements_
-
-  override def statements(fromIndex: Integer, toIndex: Integer): Seq[Statement] = statements_.slice(fromIndex, toIndex).toSeq
+  override def statements: IndexedSeq[Statement] = statements_
 
   override def interrupt(): Future[Unit] = {
     stop()

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

@@ -44,8 +44,8 @@ class WebSession(val id: Int,
   private[this] var _lastActivity = Long.MaxValue
   private[this] var _url: Option[URL] = None
 
-  private[this] var executedStatements = 0
-  private[this] var statements_ = new ArrayBuffer[Statement]
+  private[this] var _executedStatements = 0
+  private[this] var _statements = IndexedSeq[Statement]()
 
   override def url: Option[URL] = _url
 
@@ -80,22 +80,16 @@ class WebSession(val id: Int,
         }
       }
 
-      var statement = new Statement(executedStatements, content, future)
+      var statement = new Statement(_executedStatements, content, future)
 
-      executedStatements += 1
-      statements_ += statement
+      _executedStatements += 1
+      _statements = _statements :+ statement
 
       statement
     }
   }
 
-  override def statement(statementId: Int): Option[Statement] = statements_.lift(statementId)
-
-  override def statements(): Seq[Statement] = statements_.toSeq
-
-  override def statements(fromIndex: Integer, toIndex: Integer): Seq[Statement] = {
-    statements_.slice(fromIndex, toIndex).toSeq
-  }
+  override def statements: IndexedSeq[Statement] = _statements
 
   override def interrupt(): Future[Unit] = {
     stop()

+ 2 - 2
apps/spark/java/livy-server/src/test/scala/com/cloudera/hue/livy/server/BaseSessionSpec.scala

@@ -58,7 +58,7 @@ abstract class BaseSessionSpec extends FunSpec with Matchers with BeforeAndAfter
     it("should execute `1 + 2` == 3") {
       session.waitForStateChange(Starting(), Duration(30, TimeUnit.SECONDS))
       val stmt = session.executeStatement(ExecuteRequest("1 + 2"))
-      val result = Await.result(stmt.output, Duration.Inf)
+      val result = Await.result(stmt.output(), Duration.Inf)
 
       val expectedResult = Extraction.decompose(Map(
         "status" -> "ok",
@@ -74,7 +74,7 @@ abstract class BaseSessionSpec extends FunSpec with Matchers with BeforeAndAfter
     it("should report an error if accessing an unknown variable") {
       session.waitForStateChange(Starting(), Duration(30, TimeUnit.SECONDS))
       val stmt = session.executeStatement(ExecuteRequest("x"))
-      val result = Await.result(stmt.output, Duration.Inf)
+      val result = Await.result(stmt.output(), Duration.Inf)
       val expectedResult = Extraction.decompose(Map(
         "status" -> "error",
         "execution_count" -> 0,

+ 4 - 10
apps/spark/java/livy-server/src/test/scala/com/cloudera/hue/livy/server/SessionServletSpec.scala

@@ -24,13 +24,11 @@ import java.util.concurrent.atomic.AtomicInteger
 import com.cloudera.hue.livy.msgs.ExecuteRequest
 import com.cloudera.hue.livy.server.sessions._
 import com.cloudera.hue.livy.sessions._
-import org.json4s.JValue
-import org.json4s.JsonAST.{JObject, JArray}
+import org.json4s.JsonAST.{JArray, JObject}
 import org.json4s.jackson.JsonMethods._
 import org.scalatest.FunSpecLike
 import org.scalatra.test.scalatest.ScalatraSuite
 
-import scala.collection.mutable.ArrayBuffer
 import scala.concurrent.Future
 
 class SessionServletSpec extends ScalatraSuite with FunSpecLike {
@@ -39,7 +37,7 @@ class SessionServletSpec extends ScalatraSuite with FunSpecLike {
     var _state: State = Idle()
 
     var _idCounter = new AtomicInteger()
-    var _statements: ArrayBuffer[Statement] = ArrayBuffer()
+    var _statements = IndexedSeq[Statement]()
 
     override def kind: Kind = Spark()
 
@@ -58,7 +56,7 @@ class SessionServletSpec extends ScalatraSuite with FunSpecLike {
         executeRequest,
         Future.successful(JObject()))
 
-      _statements += statement
+      _statements :+= statement
 
       statement
     }
@@ -67,11 +65,7 @@ class SessionServletSpec extends ScalatraSuite with FunSpecLike {
 
     override def url: Option[URL] = ???
 
-    override def statement(statementId: Int): Option[Statement] = ???
-
-    override def statements(): Seq[Statement] = _statements
-
-    override def statements(fromIndex: Integer, toIndex: Integer): Seq[Statement] = ???
+    override def statements: IndexedSeq[Statement] = _statements
 
     override def interrupt(): Future[Unit] = ???
   }

+ 78 - 0
apps/spark/java/livy-server/src/test/scala/com/cloudera/hue/livy/server/StatementSpec.scala

@@ -0,0 +1,78 @@
+/*
+ * Licensed to Cloudera, Inc. under one
+      val statement = Statement(
+        0,
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  Cloudera, Inc. licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.cloudera.hue.livy.server
+
+import com.cloudera.hue.livy.msgs.ExecuteRequest
+import com.cloudera.hue.livy.server.sessions.Statement
+import org.json4s.JsonAST.{JArray, JString}
+import org.json4s.{DefaultFormats, Extraction}
+import org.scalatest.{Matchers, FunSpec}
+
+import scala.concurrent.duration.Duration
+import scala.concurrent.{Await, Future}
+
+class StatementSpec extends FunSpec with Matchers {
+
+  implicit val formats = DefaultFormats
+
+  describe("A statement") {
+    it("should support paging through text/plain data") {
+      val lines = List("1", "2", "3", "4", "5")
+      val rep = Extraction.decompose(Map(
+        "status" -> "ok",
+        "execution_count" -> 0,
+        "data" -> Map(
+          "text/plain" -> lines.mkString("\n")
+        )
+      ))
+      val stmt = new Statement(0, ExecuteRequest(""), Future.successful(rep))
+      var output = Await.result(stmt.output(), Duration.Inf)
+      output \ "data" \ "text/plain" should equal (JString(lines.mkString("\n")))
+
+      output = Await.result(stmt.output(Some(2)), Duration.Inf)
+      output \ "data" \ "text/plain" should equal (JString(lines.slice(2, lines.length).mkString("\n")))
+
+      output = Await.result(stmt.output(Some(2), Some(1)), Duration.Inf)
+      output \ "data" \ "text/plain" should equal (JString(lines.slice(2, 3).mkString("\n")))
+    }
+
+    it("should support paging through application/json arrays") {
+      val lines = List("1", "2", "3", "4")
+      val rep = Extraction.decompose(Map(
+        "status" -> "ok",
+        "execution_count" -> 0,
+        "data" -> Map(
+          "appliaction/json" -> List(1, 2, 3, 4)
+        )
+      ))
+      val stmt = new Statement(0, ExecuteRequest(""), Future.successful(rep))
+      var output = Await.result(stmt.output(), Duration.Inf)
+      (output \ "data" \ "text/plain").extract[List[Int]] should equal (List(1, 2, 3, 4))
+
+      output = Await.result(stmt.output(Some(2)), Duration.Inf)
+      (output \ "data" \ "text/plain").extract[List[Int]] should equal (List(3, 4))
+
+      output = Await.result(stmt.output(Some(2), Some(1)), Duration.Inf)
+      (output \ "data" \ "text/plain").extract[List[Int]] should equal (List(3))
+    }
+  }
+
+}