Bläddra i källkod

[livy] Change session/statement creation to return a 201 CREATED

Erick Tryzelaar 10 år sedan
förälder
incheckning
0e978cd

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

@@ -6,8 +6,8 @@ import org.json4s.{DefaultFormats, Formats, MappingException}
 import org.scalatra._
 import org.scalatra.json.JacksonJsonSupport
 
-import scala.concurrent.duration.Duration
-import scala.concurrent.{Future, Await, ExecutionContext, ExecutionContextExecutor}
+import scala.concurrent._
+import scala.concurrent.duration._
 
 object WebApp {
   case class CreateSessionRequest(lang: String)
@@ -34,6 +34,13 @@ class WebApp(sessionManager: SessionManager)
     sessionManager.getSessionIds
   }
 
+  val getSession = get("/sessions/:sessionId") {
+    sessionManager.get(params("sessionId")) match {
+      case Some(session) => formatSession(session)
+      case None => NotFound("Session not found")
+    }
+  }
+
   post("/sessions") {
     val createSessionRequest = parsedBody.extract[CreateSessionRequest]
 
@@ -43,7 +50,9 @@ class WebApp(sessionManager: SessionManager)
     }
 
     val rep = sessionFuture.map {
-      case session => formatSession(session)
+      case session =>
+        Created(formatSession(session),
+          headers = Map("Location" -> url(getSession, "sessionId" -> session.id.toString)))
     }
 
     // FIXME: this is silently eating exceptions.
@@ -51,13 +60,6 @@ class WebApp(sessionManager: SessionManager)
     Await.result(rep, Duration.Inf)
   }
 
-  get("/sessions/:sessionId") {
-    sessionManager.get(params("sessionId")) match {
-      case Some(session) => formatSession(session)
-      case None => NotFound("Session not found")
-    }
-  }
-
   post("/sessions/:sessionId/stop") {
     sessionManager.get(params("sessionId")) match {
       case Some(session) =>
@@ -101,35 +103,34 @@ class WebApp(sessionManager: SessionManager)
     }
   }
 
-  post("/sessions/:sessionId/statements") {
-    val req = parsedBody.extract[ExecuteStatementRequest]
-
+  val getStatement = get("/sessions/:sessionId/statements/:statementId") {
     sessionManager.get(params("sessionId")) match {
       case Some(session) =>
-        Future {
-          val statement: Statement = session.executeStatement(req.statement)
-
-          // FIXME: this is silently eating exceptions.
-          //new AsyncResult() { val is = statement }
-          Await.result(statement.output, Duration.Inf)
+        session.statement(params("statementId").toInt) match {
+          case Some(statement) => formatStatement(statement)
+          case None => NotFound("Statement not found")
         }
-
-        Accepted()
       case None => NotFound("Session not found")
     }
   }
 
-  get("/sessions/:sessionId/statements/:statementId") {
+  post("/sessions/:sessionId/statements") {
+    val req = parsedBody.extract[ExecuteStatementRequest]
+
     sessionManager.get(params("sessionId")) match {
       case Some(session) =>
-        session.statement(params("statementId").toInt) match {
-          case Some(statement) => formatStatement(statement)
-          case None => NotFound("Statement not found")
-        }
+        val statement = session.executeStatement(req.statement)
+
+        Created(formatStatement(statement),
+          headers = Map(
+            "Location" -> url(getStatement,
+              "sessionId" -> session.id.toString,
+              "statementId" -> statement.id.toString)))
       case None => NotFound("Session not found")
     }
   }
 
+
   error {
     case e: JsonParseException => halt(400, e.getMessage)
     case e: MappingException => halt(400, e.getMessage)
@@ -146,10 +147,17 @@ class WebApp(sessionManager: SessionManager)
   }
 
   private def formatStatement(statement: Statement) = {
+    // Take a couple milliseconds to see if the statement has finished.
+    val output = try {
+      Await.result(statement.output, 10 milliseconds)
+    } catch {
+      case _: TimeoutException => null
+    }
+
     Map(
       "id" -> statement.id,
       "state" -> statement.state.getClass.getSimpleName.toLowerCase,
-      "output" -> statement.output
+      "output" -> output
     )
   }
 }

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

@@ -26,13 +26,13 @@ abstract class SparkWebSession(val id: String, hostname: String, port: Int) exte
 
   override def state: State = _state
 
-  override def executeStatement(statement: String): Statement = {
+  override def executeStatement(input: String): Statement = {
     ensureIdle {
       _state = Busy()
       touchLastActivity()
 
       var req = (svc / "statements").setContentType("application/json", "UTF-8")
-      req = req << write(ExecuteRequest(statement))
+      req = req << write(ExecuteRequest(input))
 
       val future = Http(req OK as.json4s.Json).map { case (resp) =>
         synchronized {
@@ -41,11 +41,12 @@ abstract class SparkWebSession(val id: String, hostname: String, port: Int) exte
         }
       }
 
+      var statement = new Statement(executedStatements, input, future)
+
       executedStatements += 1
-      var statement_ = new Statement(executedStatements, statement, future)
-      statements_ += statement_
+      statements_ += statement
 
-      statement_
+      statement
     }
   }