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 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]
   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
     val sessionId = params("sessionId").toInt
 
 
     sessionManager.get(sessionId) match {
     sessionManager.get(sessionId) match {
+      case None => NotFound("Session not found")
       case Some(session: Session) =>
       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(
         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 sessionId = params("sessionId").toInt
     val statementId = params("statementId").toInt
     val statementId = params("statementId").toInt
 
 
+    val from = params.get("from").map(_.toInt)
+    val size = params.get("size").map(_.toInt)
+
     sessionManager.get(sessionId) match {
     sessionManager.get(sessionId) match {
+      case None => NotFound("Session not found")
       case Some(session) =>
       case Some(session) =>
-        session.statement(statementId) match {
-          case Some(statement) => statement
+        session.statements.lift(statementId) match {
           case None => NotFound("Statement not found")
           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)
   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 => ( {
   case object SessionSerializer extends CustomSerializer[Session](implicit formats => ( {
     // We don't support deserialization.
     // We don't support deserialization.
     PartialFunction.empty
     PartialFunction.empty
   }, {
   }, {
     case session: Session =>
     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
     PartialFunction.empty
   }, {
   }, {
     case statement: Statement =>
     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 => ( {
   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 com.cloudera.hue.livy.msgs.ExecuteRequest
 import org.json4s.JValue
 import org.json4s.JValue
+import org.json4s.JsonAST.{JArray, JObject, JField, JString}
 
 
 import scala.concurrent.{ExecutionContext, ExecutionContextExecutor, Future}
 import scala.concurrent.{ExecutionContext, ExecutionContextExecutor, Future}
 import scala.util.{Failure, Success}
 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._
   import Statement._
 
 
   protected implicit def executor: ExecutionContextExecutor = ExecutionContext.global
   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
   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 Success(_) => _state = Available()
     case Failure(_) => _state = Error()
     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
     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] = {
   override def interrupt(): Future[Unit] = {
     stop()
     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 _lastActivity = Long.MaxValue
   private[this] var _url: Option[URL] = None
   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
   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
       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] = {
   override def interrupt(): Future[Unit] = {
     stop()
     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") {
     it("should execute `1 + 2` == 3") {
       session.waitForStateChange(Starting(), Duration(30, TimeUnit.SECONDS))
       session.waitForStateChange(Starting(), Duration(30, TimeUnit.SECONDS))
       val stmt = session.executeStatement(ExecuteRequest("1 + 2"))
       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(
       val expectedResult = Extraction.decompose(Map(
         "status" -> "ok",
         "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") {
     it("should report an error if accessing an unknown variable") {
       session.waitForStateChange(Starting(), Duration(30, TimeUnit.SECONDS))
       session.waitForStateChange(Starting(), Duration(30, TimeUnit.SECONDS))
       val stmt = session.executeStatement(ExecuteRequest("x"))
       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(
       val expectedResult = Extraction.decompose(Map(
         "status" -> "error",
         "status" -> "error",
         "execution_count" -> 0,
         "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.msgs.ExecuteRequest
 import com.cloudera.hue.livy.server.sessions._
 import com.cloudera.hue.livy.server.sessions._
 import com.cloudera.hue.livy.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.json4s.jackson.JsonMethods._
 import org.scalatest.FunSpecLike
 import org.scalatest.FunSpecLike
 import org.scalatra.test.scalatest.ScalatraSuite
 import org.scalatra.test.scalatest.ScalatraSuite
 
 
-import scala.collection.mutable.ArrayBuffer
 import scala.concurrent.Future
 import scala.concurrent.Future
 
 
 class SessionServletSpec extends ScalatraSuite with FunSpecLike {
 class SessionServletSpec extends ScalatraSuite with FunSpecLike {
@@ -39,7 +37,7 @@ class SessionServletSpec extends ScalatraSuite with FunSpecLike {
     var _state: State = Idle()
     var _state: State = Idle()
 
 
     var _idCounter = new AtomicInteger()
     var _idCounter = new AtomicInteger()
-    var _statements: ArrayBuffer[Statement] = ArrayBuffer()
+    var _statements = IndexedSeq[Statement]()
 
 
     override def kind: Kind = Spark()
     override def kind: Kind = Spark()
 
 
@@ -58,7 +56,7 @@ class SessionServletSpec extends ScalatraSuite with FunSpecLike {
         executeRequest,
         executeRequest,
         Future.successful(JObject()))
         Future.successful(JObject()))
 
 
-      _statements += statement
+      _statements :+= statement
 
 
       statement
       statement
     }
     }
@@ -67,11 +65,7 @@ class SessionServletSpec extends ScalatraSuite with FunSpecLike {
 
 
     override def url: Option[URL] = ???
     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] = ???
     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))
+    }
+  }
+
+}