org.apache.hadoop.hive.ql.session.SessionState Scala Examples

package com.webank.wedatasphere.linkis.engine.hive.executor

import java.io.PrintStream

import com.webank.wedatasphere.linkis.engine.execute.{EngineExecutor, EngineExecutorFactory}
import com.webank.wedatasphere.linkis.engine.hive.common.HiveUtils
import com.webank.wedatasphere.linkis.engine.hive.exception.HiveSessionStartFailedException
import com.webank.wedatasphere.linkis.server.JMap
import org.apache.hadoop.hive.conf.HiveConf
import org.apache.hadoop.hive.ql.Driver
import org.apache.hadoop.hive.ql.session.SessionState
import org.apache.hadoop.security.UserGroupInformation
import org.slf4j.LoggerFactory
import org.springframework.stereotype.Component


@Component
class HiveEngineExecutorFactory extends EngineExecutorFactory {
  private val logger = LoggerFactory.getLogger(getClass)
  private val HIVE_QUEUE_NAME:String = "mapreduce.job.queuename"
  private val BDP_QUEUE_NAME:String = "wds.linkis.yarnqueue"
  override def createExecutor(options: JMap[String, String]): EngineExecutor = {
    val hiveConf:HiveConf = HiveUtils.getHiveConf
    hiveConf.setVar(HiveConf.ConfVars.HIVEJAR, HiveUtils.jarOfClass(classOf[Driver])
      .getOrElse(throw HiveSessionStartFailedException(40012 ,"cannot find hive-exec.jar, start session failed!")))
    import scala.collection.JavaConversions._
    options.foreach{ case(k,v) => logger.info(s"key is $k, value is $v")}
    options.filter{case (k,v) => k.startsWith("hive.") || k.startsWith("mapreduce.") || k.startsWith("wds.linkis.")}.foreach{case(k, v) =>
      logger.info(s"key is $k, value is $v")
      if (BDP_QUEUE_NAME.equals(k)) hiveConf.set(HIVE_QUEUE_NAME, v) else hiveConf.set(k, v)}
    val sessionState:SessionState = new SessionState(hiveConf)
    sessionState.out = new PrintStream(System.out, true, "utf-8")
    sessionState.info = new PrintStream(System.out, true, "utf-8")
    sessionState.err = new PrintStream(System.out, true, "utf-8")
    SessionState.start(sessionState)
    val ugi = UserGroupInformation.getCurrentUser
    new HiveEngineExecutor(5000, sessionState, ugi, hiveConf)
  }

} 

package org.apache.spark.sql.hive.thriftserver

import org.apache.hadoop.hive.cli.CliSessionState
import org.apache.hadoop.hive.conf.HiveConf
import org.apache.hadoop.hive.ql.session.SessionState

import org.apache.spark.{SparkConf, SparkFunSuite}
import org.apache.spark.deploy.SparkHadoopUtil
import org.apache.spark.sql.hive.HiveUtils

class HiveCliSessionStateSuite extends SparkFunSuite {

  def withSessionClear(f: () => Unit): Unit = {
    try f finally SessionState.detachSession()
  }

  test("CliSessionState will be reused") {
    withSessionClear { () =>
      val hiveConf = new HiveConf(classOf[SessionState])
      HiveUtils.newTemporaryConfiguration(useInMemoryDerby = false).foreach {
        case (key, value) => hiveConf.set(key, value)
      }
      val sessionState: SessionState = new CliSessionState(hiveConf)
      SessionState.start(sessionState)
      val s1 = SessionState.get
      val sparkConf = new SparkConf()
      val hadoopConf = SparkHadoopUtil.get.newConfiguration(sparkConf)
      val s2 = HiveUtils.newClientForMetadata(sparkConf, hadoopConf).getState
      assert(s1 === s2)
      assert(s2.isInstanceOf[CliSessionState])
    }
  }

  test("SessionState will not be reused") {
    withSessionClear { () =>
      val sparkConf = new SparkConf()
      val hadoopConf = SparkHadoopUtil.get.newConfiguration(sparkConf)
      HiveUtils.newTemporaryConfiguration(useInMemoryDerby = false).foreach {
        case (key, value) => hadoopConf.set(key, value)
      }
      val hiveClient = HiveUtils.newClientForMetadata(sparkConf, hadoopConf)
      val s1 = hiveClient.getState
      val s2 = hiveClient.newSession().getState
      assert(s1 !== s2)
    }
  }
} 

package org.apache.spark.sql.hive.client

import java.util.{List => JList}

import com.githup.yaooqinn.spark.authorizer.Logging
import org.apache.hadoop.hive.ql.security.authorization.plugin._
import org.apache.hadoop.hive.ql.session.SessionState
import org.apache.hadoop.security.UserGroupInformation

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.hive.{AuthzUtils, HiveExternalCatalog}
import org.apache.spark.sql.internal.NonClosableMutableURLClassLoader


object AuthzImpl extends Logging {
  def checkPrivileges(
      spark: SparkSession,
      hiveOpType: HiveOperationType,
      inputObjs: JList[HivePrivilegeObject],
      outputObjs: JList[HivePrivilegeObject],
      context: HiveAuthzContext): Unit = {
    val client = spark.sharedState
      .externalCatalog.asInstanceOf[HiveExternalCatalog]
      .client
    val clientImpl = try {
      client.asInstanceOf[HiveClientImpl]
    } catch {
      case _: ClassCastException =>
        val clientLoader =
          AuthzUtils.getFieldVal(client, "clientLoader").asInstanceOf[IsolatedClientLoader]
        AuthzUtils.setFieldVal(clientLoader, "isolationOn", false)
        AuthzUtils.setFieldVal(clientLoader,
          "classLoader", new NonClosableMutableURLClassLoader(clientLoader.baseClassLoader))
        clientLoader.cachedHive = null
        val newClient = clientLoader.createClient()
        AuthzUtils.setFieldVal(
          spark.sharedState.externalCatalog.asInstanceOf[HiveExternalCatalog],
          "client",
          newClient)
        newClient.asInstanceOf[HiveClientImpl]
    }

    val state = clientImpl.state
    SessionState.setCurrentSessionState(state)
    val user = UserGroupInformation.getCurrentUser.getShortUserName
    if (state.getAuthenticator.getUserName != user) {
      val hiveConf = state.getConf
      val newState = new SessionState(hiveConf, user)
      SessionState.start(newState)
      AuthzUtils.setFieldVal(clientImpl, "state", newState)
    }

    val authz = clientImpl.state.getAuthorizerV2
    clientImpl.withHiveState {
      if (authz != null) {
        try {
          authz.checkPrivileges(hiveOpType, inputObjs, outputObjs, context)
        } catch {
          case hae: HiveAccessControlException =>
            error(
              s"""
                 |+===============================+
                 ||Spark SQL Authorization Failure|
                 ||-------------------------------|
                 ||${hae.getMessage}
                 ||-------------------------------|
                 ||Spark SQL Authorization Failure|
                 |+===============================+
               """.stripMargin)
            throw hae
          case e: Exception => throw e
        }
      } else {
        warn("Authorizer V2 not configured. Skipping privilege checking")
      }
    }
  }
} 

package yaooqinn.kyuubi.utils

import org.apache.hadoop.fs.FileSystem
import org.apache.hadoop.hive.conf.HiveConf
import org.apache.hadoop.hive.ql.session.SessionState
import org.apache.hadoop.security.UserGroupInformation
import org.apache.kyuubi.Logging
import org.apache.spark.{KyuubiSparkUtil, SparkConf}

object KyuubiHiveUtil extends Logging {

  private val HIVE_PREFIX = "hive."
  private val METASTORE_PREFIX = "metastore."

  val URIS: String = HIVE_PREFIX + METASTORE_PREFIX + "uris"
  val METASTORE_PRINCIPAL: String = HIVE_PREFIX + METASTORE_PREFIX + "kerberos.principal"

  def hiveConf(conf: SparkConf): HiveConf = {
    val hadoopConf = KyuubiSparkUtil.newConfiguration(conf)
    new HiveConf(hadoopConf, classOf[HiveConf])
  }

  def addDelegationTokensToHiveState(ugi: UserGroupInformation): Unit = {
    val state = SessionState.get
    if (state != null) {
      addDelegationTokensToHiveState(state, ugi)
    }
  }

  def addDelegationTokensToHiveState(state: SessionState, ugi: UserGroupInformation): Unit = {
    state.getHdfsEncryptionShim match {
      case shim: org.apache.hadoop.hive.shims.Hadoop23Shims#HdfsEncryptionShim =>
        try {
          val hdfsAdmin = ReflectUtils.getFieldValue(shim, "hdfsAdmin")
          val dfs = ReflectUtils.getFieldValue(hdfsAdmin, "dfs")
          dfs.asInstanceOf[FileSystem].addDelegationTokens(ugi.getUserName, ugi.getCredentials)
        } catch {
          case e: Exception =>
            error("Failed add delegation token to hive session state", e)
        }
      case _ =>
    }
  }
} 

package yaooqinn.kyuubi.operation

import java.io.{File, IOException}

import scala.util.Try

import org.apache.hadoop.fs.Path
import org.apache.hadoop.hdfs.{HdfsConfiguration, MiniDFSCluster}
import org.apache.hadoop.hive.conf.HiveConf
import org.apache.hadoop.hive.ql.session.SessionState
import org.apache.hadoop.security.UserGroupInformation
import org.apache.spark.sql.catalyst.catalog.FunctionResource
import org.apache.spark.sql.execution.SparkSqlParser
import org.apache.spark.sql.internal.SQLConf
import org.mockito.Mockito.when

import yaooqinn.kyuubi.operation.statement.ExecuteStatementInClientMode
import yaooqinn.kyuubi.utils.{KyuubiHiveUtil, ReflectUtils}

class ExecuteStatementInClientModeWithHDFSSuite extends ExecuteStatementInClientModeSuite {
  val hdfsConf = new HdfsConfiguration
  hdfsConf.set("fs.hdfs.impl.disable.cache", "true")
  var cluster: MiniDFSCluster = new MiniDFSCluster.Builder(hdfsConf).build()
  cluster.waitClusterUp()
  val fs = cluster.getFileSystem
  val homeDirectory: Path = fs.getHomeDirectory
  private val fileName = "example-1.0.0-SNAPSHOT.jar"
  private val remoteUDFFile = new Path(homeDirectory, fileName)

  override def beforeAll(): Unit = {
    val file = new File(this.getClass.getProtectionDomain.getCodeSource.getLocation + fileName)
    val localUDFFile = new Path(file.getPath)
    fs.copyFromLocalFile(localUDFFile, remoteUDFFile)
    super.beforeAll()
  }

  override def afterAll(): Unit = {
    fs.delete(remoteUDFFile, true)
    fs.close()
    cluster.shutdown()
    super.afterAll()
  }

  test("transform logical plan") {
    val op = sessionMgr.getOperationMgr.newExecuteStatementOperation(session, statement)
      .asInstanceOf[ExecuteStatementInClientMode]
    val parser = new SparkSqlParser(new SQLConf)
    val plan0 = parser.parsePlan(
      s"create temporary function a as 'a.b.c' using file '$remoteUDFFile'")
    val plan1 = op.transform(plan0)
    assert(plan0 === plan1)
    assert(
      ReflectUtils.getFieldValue(plan1, "resources").asInstanceOf[Seq[FunctionResource]].isEmpty)

    val plan2 = parser.parsePlan(
      s"create temporary function a as 'a.b.c' using jar '$remoteUDFFile'")
    val plan3 = op.transform(plan2)
    assert(plan3 === plan2)
    assert(
      ReflectUtils.getFieldValue(plan3, "resources").asInstanceOf[Seq[FunctionResource]].isEmpty)
  }

  test("add delegation token with hive session state, hdfs") {
    val hiveConf = new HiveConf(hdfsConf, classOf[HiveConf])
    val state = new SessionState(hiveConf)
    assert(Try {
      KyuubiHiveUtil.addDelegationTokensToHiveState(state, UserGroupInformation.getCurrentUser)
    }.isSuccess)

    val mockuser = mock[UserGroupInformation]
    when(mockuser.getUserName).thenThrow(classOf[IOException])
    KyuubiHiveUtil.addDelegationTokensToHiveState(state, mockuser)
  }
} 

package yaooqinn.kyuubi.utils

import scala.util.Try

import org.apache.hadoop.hive.ql.session.SessionState
import org.apache.hadoop.security.UserGroupInformation
import org.apache.spark.{SparkConf, SparkFunSuite}

class KyuubiHiveUtilSuite extends SparkFunSuite {
  private val user = UserGroupInformation.getCurrentUser

  test("hive conf") {
    val uris = "thrift://yaooqinn.kyuubi"
    val conf = new SparkConf()
      .set("spark.hadoop.hive.metastore.uris", uris)
    val hiveConf = KyuubiHiveUtil.hiveConf(conf)
    assert(hiveConf.get(KyuubiHiveUtil.URIS) === uris)
  }

  test("testURIS") {
    assert(KyuubiHiveUtil.URIS === "hive.metastore.uris")
  }

  test("metastore principal") {
    assert(KyuubiHiveUtil.METASTORE_PRINCIPAL === "hive.metastore.kerberos.principal")

  }

  test("add delegation tokens without hive session state ") {
    assert(Try {KyuubiHiveUtil.addDelegationTokensToHiveState(user)}.isSuccess)

  }

  test("add delegation token with hive session state, local fs") {
    val state = new SessionState(KyuubiHiveUtil.hiveConf(new SparkConf()))
    assert(Try {KyuubiHiveUtil.addDelegationTokensToHiveState(state, user) }.isSuccess)
  }

} 

package org.apache.spark.sql.hive.thriftserver

import org.apache.hadoop.hive.cli.CliSessionState
import org.apache.hadoop.hive.conf.HiveConf
import org.apache.hadoop.hive.ql.session.SessionState

import org.apache.spark.{SparkConf, SparkFunSuite}
import org.apache.spark.deploy.SparkHadoopUtil
import org.apache.spark.sql.hive.HiveUtils

class HiveCliSessionStateSuite extends SparkFunSuite {

  def withSessionClear(f: () => Unit): Unit = {
    try f finally SessionState.detachSession()
  }

  test("CliSessionState will be reused") {
    withSessionClear { () =>
      val hiveConf = new HiveConf(classOf[SessionState])
      HiveUtils.newTemporaryConfiguration(useInMemoryDerby = false).foreach {
        case (key, value) => hiveConf.set(key, value)
      }
      val sessionState: SessionState = new CliSessionState(hiveConf)
      SessionState.start(sessionState)
      val s1 = SessionState.get
      val sparkConf = new SparkConf()
      val hadoopConf = SparkHadoopUtil.get.newConfiguration(sparkConf)
      val s2 = HiveUtils.newClientForMetadata(sparkConf, hadoopConf).getState
      assert(s1 === s2)
      assert(s2.isInstanceOf[CliSessionState])
    }
  }

  test("SessionState will not be reused") {
    withSessionClear { () =>
      val sparkConf = new SparkConf()
      val hadoopConf = SparkHadoopUtil.get.newConfiguration(sparkConf)
      HiveUtils.newTemporaryConfiguration(useInMemoryDerby = false).foreach {
        case (key, value) => hadoopConf.set(key, value)
      }
      val hiveClient = HiveUtils.newClientForMetadata(sparkConf, hadoopConf)
      val s1 = hiveClient.getState
      val s2 = hiveClient.newSession().getState
      assert(s1 !== s2)
    }
  }
} 

package com.flaminem.flamy.exec.hive

import com.flaminem.flamy.conf.FlamyContext
import com.flaminem.flamy.conf.hive.ModelHiveContext
import com.flaminem.flamy.model.Variables
import com.flaminem.flamy.model.exceptions.FailedQueryException
import com.flaminem.flamy.model.files.PresetsFile
import com.flaminem.flamy.utils.logging.Logging
import com.flaminem.flamy.utils.sql.hive.StreamedResultSet
import org.apache.hadoop.hive.cli.CliDriver
import org.apache.hadoop.hive.conf.HiveConf
import org.apache.hadoop.hive.ql.Driver
import org.apache.hadoop.hive.ql.session.SessionState

class LocalHiveRunner(override val context: FlamyContext) extends HiveRunner(context: FlamyContext) with Logging {

  private val modelHiveContext: ModelHiveContext = ModelHiveContext.getHeavyContext(context)
  assert(Option(SessionState.get()).isDefined)
  private val cliDriver = new CliDriver

  context.getHadoopBinPath match {
    case None => ()
    case Some(path) => cliDriver.processCmd(s"SET hadoop.bin.path=$path")
  }

  protected override def runPreparedQuery(query: String, title: Option[String]): Int = {
    if (query.toUpperCase.startsWith("LOAD DATA") || query.toUpperCase.startsWith("EXPLAIN LOAD DATA")) {
      logger.info("Skipping query: LOAD DATA")
      0
    }
    else {
      System.out.println(query)
      cliDriver.processCmd(query)
    }
  }

  
  def getConfiguration(key: String): Option[String] = {
    val conf: HiveConf = modelHiveContext.hiveConf
    Option(conf.get(key))
  }

  override def close(): Unit = () //sessionState.close()

  override def interrupt(): Unit = () //sessionState.close()

}