Browse Source

[livy] Add support for impersonization

Erick Tryzelaar 10 years ago
parent
commit
6643fef9a8

+ 6 - 6
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}
 import scala.concurrent.{ExecutionContext, Future}
 
 
 trait SessionFactory {
 trait SessionFactory {
-  def createSession(kind: Session.Kind): Future[Session]
+  def createSession(kind: Session.Kind, proxyUser: Option[String] = None): Future[Session]
 
 
   def close(): Unit = {}
   def close(): Unit = {}
 }
 }
@@ -18,7 +18,7 @@ class ThreadSessionFactory(livyConf: LivyConf) extends SessionFactory {
 
 
   implicit def executor: ExecutionContext = ExecutionContext.global
   implicit def executor: ExecutionContext = ExecutionContext.global
 
 
-  override def createSession(kind: Session.Kind): Future[Session] = {
+  override def createSession(kind: Session.Kind, proxyUser: Option[String] = None): Future[Session] = {
     Future {
     Future {
       val id = UUID.randomUUID().toString
       val id = UUID.randomUUID().toString
       ThreadSession.create(id, kind)
       ThreadSession.create(id, kind)
@@ -30,10 +30,10 @@ class ProcessSessionFactory(livyConf: LivyConf) extends SessionFactory {
 
 
   implicit def executor: ExecutionContext = ExecutionContext.global
   implicit def executor: ExecutionContext = ExecutionContext.global
 
 
-  override def createSession(kind: Session.Kind): Future[Session] = {
+  override def createSession(kind: Session.Kind, proxyUser: Option[String] = None): Future[Session] = {
     Future {
     Future {
       val id = UUID.randomUUID().toString
       val id = UUID.randomUUID().toString
-      ProcessSession.create(livyConf, id, kind)
+      ProcessSession.create(livyConf, id, kind, proxyUser)
     }
     }
   }
   }
 }
 }
@@ -42,9 +42,9 @@ class YarnSessionFactory(livyConf: LivyConf) extends SessionFactory {
 
 
   val client = new Client(livyConf)
   val client = new Client(livyConf)
 
 
-  override def createSession(kind: Session.Kind): Future[Session] = {
+  override def createSession(kind: Session.Kind, proxyUser: Option[String] = None): Future[Session] = {
     val id = UUID.randomUUID().toString
     val id = UUID.randomUUID().toString
-    YarnSession.create(client, id, kind)
+    YarnSession.create(client, id, kind, proxyUser)
   }
   }
 
 
   override def close(): Unit = {
   override def close(): Unit = {

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

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

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

@@ -47,7 +47,7 @@ class WebApp(sessionManager: SessionManager)
 
 
   post("/sessions") {
   post("/sessions") {
     val createSessionRequest = parsedBody.extract[CreateSessionRequest]
     val createSessionRequest = parsedBody.extract[CreateSessionRequest]
-    val sessionFuture = sessionManager.createSession(createSessionRequest.lang)
+    val sessionFuture = sessionManager.createSession(createSessionRequest.lang, createSessionRequest.proxyUser)
 
 
     val rep = sessionFuture.map { case session =>
     val rep = sessionFuture.map { case session =>
       Created(session,
       Created(session,
@@ -179,7 +179,8 @@ private object Serializers {
 
 
       ("id", session.id) ~
       ("id", session.id) ~
       ("state", serializeSessionState(session.state)) ~
       ("state", serializeSessionState(session.state)) ~
-      ("kind", serializeSessionKind(session.kind))
+      ("kind", serializeSessionKind(session.kind)) ~
+      ("proxyUser", session.proxyUser)
   }
   }
     )
     )
   )
   )

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

@@ -10,7 +10,6 @@ import scala.collection.JavaConversions._
 import scala.collection.mutable.ArrayBuffer
 import scala.collection.mutable.ArrayBuffer
 import scala.concurrent.Future
 import scala.concurrent.Future
 import scala.io.Source
 import scala.io.Source
-import scala.util.control.Breaks._
 
 
 object ProcessSession extends Logging {
 object ProcessSession extends Logging {
 
 
@@ -18,13 +17,13 @@ object ProcessSession extends Logging {
   val CONF_LIVY_REPL_CALLBACK_URL = "livy.repl.callback-url"
   val CONF_LIVY_REPL_CALLBACK_URL = "livy.repl.callback-url"
   val CONF_LIVY_REPL_DRIVER_CLASS_PATH = "livy.repl.driverClassPath"
   val CONF_LIVY_REPL_DRIVER_CLASS_PATH = "livy.repl.driverClassPath"
 
 
-  def create(livyConf: LivyConf, id: String, kind: Session.Kind): Session = {
-    val process = startProcess(livyConf, id, kind)
-    new ProcessSession(id, kind, process)
+  def create(livyConf: LivyConf, id: String, kind: Session.Kind, proxyUser: Option[String] = None): Session = {
+    val process = startProcess(livyConf, id, kind, proxyUser)
+    new ProcessSession(id, kind, proxyUser, process)
   }
   }
 
 
   // Loop until we've started a process with a valid port.
   // Loop until we've started a process with a valid port.
-  private def startProcess(livyConf: LivyConf, id: String, kind: Session.Kind): Process = {
+  private def startProcess(livyConf: LivyConf, id: String, kind: Session.Kind, proxyUser: Option[String]): Process = {
     val args = ArrayBuffer(
     val args = ArrayBuffer(
       "spark-submit",
       "spark-submit",
       "--class", "com.cloudera.hue.livy.repl.Main"
       "--class", "com.cloudera.hue.livy.repl.Main"
@@ -40,6 +39,11 @@ object ProcessSession extends Logging {
       args += extraClassPath
       args += extraClassPath
     }
     }
 
 
+    proxyUser.foreach { case user =>
+      args += "--proxy-user"
+      args += user
+    }
+
     args += livyJar(livyConf)
     args += livyJar(livyConf)
     args += kind.toString
     args += kind.toString
 
 
@@ -64,7 +68,10 @@ object ProcessSession extends Logging {
   }
   }
 }
 }
 
 
-private class ProcessSession(id: String, kind: Session.Kind, process: Process) extends WebSession(id, kind) {
+private class ProcessSession(id: String,
+                             kind: Session.Kind,
+                             proxyUser: Option[String],
+                             process: Process) extends WebSession(id, kind, proxyUser) {
 
 
   val stdoutThread = new Thread {
   val stdoutThread = new Thread {
     override def run() = {
     override def run() = {

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

@@ -59,6 +59,8 @@ trait Session {
 
 
   def kind: Kind
   def kind: Kind
 
 
+  def proxyUser: Option[String]
+
   def lastActivity: Long
   def lastActivity: Long
 
 
   def state: State
   def state: State

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

@@ -27,13 +27,17 @@ object ThreadSession {
   }
   }
 }
 }
 
 
-private class ThreadSession(val id: String, val kind: Session.Kind, 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
   protected implicit def executor: ExecutionContextExecutor = ExecutionContext.global
 
 
   private var executedStatements = 0
   private var executedStatements = 0
   private var statements_ = new ArrayBuffer[Statement]
   private var statements_ = new ArrayBuffer[Statement]
 
 
+  override def proxyUser: Option[String] = None
+
   override def lastActivity: Long = 0
   override def lastActivity: Long = 0
 
 
   override def state: State = {
   override def state: State = {

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

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

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

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

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

@@ -10,12 +10,12 @@ import org.apache.hadoop.yarn.conf.YarnConfiguration
 import org.apache.hadoop.yarn.util.ConverterUtils
 import org.apache.hadoop.yarn.util.ConverterUtils
 
 
 import scala.annotation.tailrec
 import scala.annotation.tailrec
+import scala.collection.JavaConversions._
+import scala.collection.mutable.ArrayBuffer
 import scala.concurrent.{ExecutionContext, Future}
 import scala.concurrent.{ExecutionContext, Future}
 
 
 object Client extends Logging {
 object Client extends Logging {
-  private val LIVY_JAR = "__livy__.jar"
   private val CONF_LIVY_JAR = "livy.yarn.jar"
   private val CONF_LIVY_JAR = "livy.yarn.jar"
-  private val LOCAL_SCHEME = "local"
   private lazy val regex = """Application report for (\w+)""".r.unanchored
   private lazy val regex = """Application report for (\w+)""".r.unanchored
 
 
   private def livyJar(conf: LivyConf) = {
   private def livyJar(conf: LivyConf) = {
@@ -39,18 +39,29 @@ class Client(livyConf: LivyConf) extends Logging {
   yarnClient.init(yarnConf)
   yarnClient.init(yarnConf)
   yarnClient.start()
   yarnClient.start()
 
 
-  def submitApplication(id: String, kind: String, callbackUrl: String): Future[Job] = {
+  def submitApplication(id: String,
+                        kind: String,
+                        proxyUser: Option[String],
+                        callbackUrl: String): Future[Job] = {
     val url = f"$callbackUrl/sessions/$id/callback"
     val url = f"$callbackUrl/sessions/$id/callback"
 
 
-    val builder: ProcessBuilder = new ProcessBuilder(
+    val args: ArrayBuffer[String] = ArrayBuffer(
       "spark-submit",
       "spark-submit",
       "--master", "yarn-cluster",
       "--master", "yarn-cluster",
       "--class", "com.cloudera.hue.livy.repl.Main",
       "--class", "com.cloudera.hue.livy.repl.Main",
-      "--driver-java-options", f"-Dlivy.repl.callback-url=$url -Dlivy.repl.port=0",
-      livyJar(livyConf),
-      kind
+      "--driver-java-options", f"-Dlivy.repl.callback-url=$url -Dlivy.repl.port=0"
     )
     )
 
 
+    proxyUser.foreach { case user =>
+      args += "--proxy-user"
+      args += user
+    }
+
+    args += livyJar(livyConf)
+    args += kind
+
+    val builder: ProcessBuilder = new ProcessBuilder(args)
+
     builder.redirectOutput(Redirect.PIPE)
     builder.redirectOutput(Redirect.PIPE)
     builder.redirectErrorStream(true)
     builder.redirectErrorStream(true)
 
 
@@ -97,9 +108,8 @@ class Job(yarnClient: YarnClient, appId: ApplicationId) {
     while (System.currentTimeMillis() - startTimeMs < timeoutMs) {
     while (System.currentTimeMillis() - startTimeMs < timeoutMs) {
       val status = getStatus
       val status = getStatus
       status match {
       status match {
-        case SuccessfulFinish() | UnsuccessfulFinish() => {
+        case SuccessfulFinish() | UnsuccessfulFinish() =>
           return Some(status)
           return Some(status)
-        }
         case _ =>
         case _ =>
       }
       }