<commit_before>package io.eels.component.hive

import com.sksamuel.exts.Logging
import com.sksamuel.exts.io.Using
import io.eels.component.hdfs.{AclSpec, HdfsSource}
import io.eels.schema.{PartitionConstraint, PartitionEquals, Schema}
import io.eels.{FilePattern, Part, Source}
import io.eels.component.parquet.{ParquetLogMute, Predicate}
import org.apache.hadoop.fs.{FileSystem, Path}
import org.apache.hadoop.fs.permission.FsPermission
import org.apache.hadoop.hive.metastore.{IMetaStoreClient, TableType}

import scala.collection.JavaConverters._

/**
 * @param constraints optional constraits on the partition data to narrow which partitions are read
 * @param projection sets which fields are required by the caller.
 * @param predicate optional predicate which will filter rows at the read level
 *
 */
case class HiveSource(dbName: String,
                      tableName: String,
                      constraints: List[PartitionConstraint] = Nil,
                      projection: List[String] = Nil,
                      predicate: Option[Predicate] = None)
                     (implicit fs: FileSystem,
                      client: IMetaStoreClient) extends Source with Logging with Using {
    ParquetLogMute()

  implicit val conf = fs.getConf
  val ops = new HiveOps(client)

  def select(first: String, rest: String*): HiveSource = withProjection(first +: rest)
  def withProjection(first: String, rest: String*): HiveSource = withProjection(first +: rest)
  def withProjection(columns: Seq[String]): HiveSource = {
    require(columns.nonEmpty)
    copy(projection = columns.toList)
  }

  /**
    * Returns a list of all files used by this hive source.
    *
    * @param includePartitionDirs if true then the partition directories will be included
    * @param includeTableDir      if true then the main table directory will be included
    * @return paths of all files and directories
    */
  def paths(includePartitionDirs: Boolean = false, includeTableDir: Boolean = false): List[Path] = {
    val files = ops.partitions(dbName, tableName).flatMap { partition =>
      val location = partition.getSd.getLocation
      val files = FilePattern(s"$location/*").toPaths()
      if (includePartitionDirs) {
        files :+ new Path(location)
      } else {
        files
      }
    }
    if (includeTableDir) {
      val location = spec().location
      files :+ new Path(location)
    } else {
      files
    }
  }

  def setPermissions(permission: FsPermission,
                     includePartitionDirs: Boolean = false,
                     includeTableDir: Boolean = false): Unit = {
    paths(includePartitionDirs, includeTableDir).foreach(fs.setPermission(_, permission))
  }

  /**
    * Sets the acl for all files of this hive source.
    * Even if the files are not located inside the table directory, this function will find them
    * and correctly update the spec.
    *
    * @param acl the acl values to set
    */
  def setAcl(acl: AclSpec,
             includePartitionDirs: Boolean = false,
             includeTableDir: Boolean = false): Unit = {
    paths(includePartitionDirs, includeTableDir).foreach { path =>
      HdfsSource(path).setAcl(acl)
    }
  }

  // returns the permission of the table location path
  def tablePermission(): FsPermission = {
    val location = ops.location(dbName, tableName)
    fs.getFileStatus(new Path(location)).getPermission
  }

  /**
    * Returns a TableSpec which contains details of the underlying table.
    * Similar to the Table class in the Hive API but using scala friendly types.
    */
  def spec(): TableSpec = {
    val table = client.getTable(dbName, tableName)
    val tableType = TableType.values().find(_.name.toLowerCase == table.getTableType.toLowerCase)
      .getOrElse(sys.error("Hive table type is not supported by this version of hive"))
    val params = table.getParameters.asScala.toMap
    TableSpec(
      tableName,
      tableType,
      table.getSd.getLocation,
      table.getSd.getNumBuckets,
      table.getSd.getBucketCols.asScala.toList,
      params,
      table.getSd.getInputFormat,
      table.getSd.getOutputFormat,
      table.getSd.getSerdeInfo.getName,
      table.getRetention,
      table.getCreateTime,
      table.getLastAccessTime,
      table.getOwner
    )
  }

  def withPredicate(predicate: Predicate): HiveSource = copy(predicate = Some(predicate))

  def withPartitionConstraint(name: String, value: String): HiveSource = withPartitionConstraint(PartitionEquals(name, value))

  def withPartitionConstraint(expr: PartitionConstraint): HiveSource = {
    copy(constraints = constraints :+ expr)
  }

  /**
   * The returned schema should take into account:
   *
   * 1) Any projection. If a projection is set, then it should return the schema in the same order
   * as the projection. If no projection is set then the schema should be driven from the hive metastore.
   *
   * 2) Any partitions set. These should be included in the schema columns.
   */
  override def schema(): Schema = {

    // if no field names were specified, then we will return the schema as is from the hive database,
    // otherwise we will keep only the requested fields
    val schema = if (projection.isEmpty) metastoreSchema
    else {
      // remember hive is always lower case, so when comparing requested field names with
      // hive fields we need to use lower case everything. And we need to return the schema
      // in the same order as the requested projection
      val columns = projection.map { fieldName =>
        val field = metastoreSchema.fields.find { it => it.name.equals(fieldName, true) }
        field.getOrElse(sys.error(s"Requested field $fieldName does not exist in the hive schema"))
      }
      Schema(columns)
    }

    schema
  }

  // returns the full underlying schema from the metastore including partition partitionKeys
  val metastoreSchema: Schema = ops.schema(dbName, tableName)

  //def  spec(): HiveSpec = HiveSpecFn.toHiveSpec(dbName, tableName)

  def isPartitionOnlyProjection(): Boolean = {
    val partitionKeyNames = HiveTable(dbName, tableName).partitionKeys().map(_.field.name)
    projection.nonEmpty && projection.map { it => it.toLowerCase() }.forall { it => partitionKeyNames.contains(it) }
  }

  override def parts(): List[Part] = {

    val table = client.getTable(dbName, tableName)
    val dialect = io.eels.component.hive.HiveDialect(table)
    val partitionKeys = HiveTable(dbName, tableName).partitionKeys()

    // all field names from the underlying hive schema
    val fieldNames = metastoreSchema.fields.map(_.name)

    // if we requested only partition columns, then we can get this information by scanning the file system
    // to see which partitions have been created. Then the presence of files in a partition location means that
    // that partition must have data.
    if (isPartitionOnlyProjection()) {
      logger.info("Requested projection contains only partitions; reading directly from metastore")
      // we pass in the schema so we can order the results to keep them aligned with the given projection
      List(new HivePartitionPart(dbName, tableName, fieldNames, Nil, metastoreSchema, predicate, dialect))
    } else {
      val files = HiveFilesFn(table, partitionKeys.map(_.field.name), constraints)
      logger.debug(s"Found ${files.size} visible hive files from all locations for $dbName:$tableName")

      // for each seperate hive file part we must pass in the metastore schema
      files.map { case (file, spec) =>
        new HiveFilePart(dialect, file, metastoreSchema, schema(), predicate, spec.parts.toList)
      }
    }
  }
}<commit_msg>Have a withKeytab option for hive sources  #169<commit_after>package io.eels.component.hive

import com.sksamuel.exts.Logging
import com.sksamuel.exts.OptionImplicits._
import com.sksamuel.exts.io.Using
import io.eels.component.hdfs.{AclSpec, HdfsSource}
import io.eels.component.parquet.{ParquetLogMute, Predicate}
import io.eels.schema.{PartitionConstraint, PartitionEquals, Schema}
import io.eels.{FilePattern, Part, Source}
import org.apache.hadoop.fs.permission.FsPermission
import org.apache.hadoop.fs.{FileSystem, Path}
import org.apache.hadoop.hive.metastore.{IMetaStoreClient, TableType}
import org.apache.hadoop.security.UserGroupInformation

import scala.collection.JavaConverters._

/**
 * @param constraints optional constraits on the partition data to narrow which partitions are read
 * @param projection sets which fields are required by the caller.
 * @param predicate optional predicate which will filter rows at the read level
 *
 */
case class HiveSource(dbName: String,
                      tableName: String,
                      constraints: List[PartitionConstraint] = Nil,
                      projection: List[String] = Nil,
                      predicate: Option[Predicate] = None,
                      principal: Option[String] = None,
                      keytabPath: Option[java.nio.file.Path] = None)
                     (implicit fs: FileSystem,
                      client: IMetaStoreClient) extends Source with Logging with Using {
    ParquetLogMute()

  implicit val conf = fs.getConf
  val ops = new HiveOps(client)

  def select(first: String, rest: String*): HiveSource = withProjection(first +: rest)
  def withProjection(first: String, rest: String*): HiveSource = withProjection(first +: rest)
  def withProjection(columns: Seq[String]): HiveSource = {
    require(columns.nonEmpty)
    copy(projection = columns.toList)
  }

  def withPredicate(predicate: Predicate): HiveSource = copy(predicate = Some(predicate))

  def withPartitionConstraint(name: String, value: String): HiveSource = withPartitionConstraint(PartitionEquals(name, value))

  def withPartitionConstraint(expr: PartitionConstraint): HiveSource = {
    copy(constraints = constraints :+ expr)
  }

  def withKeytabFile(principal: String, keytabPath: java.nio.file.Path): HiveSource = {
    login()
    copy(principal = principal.some, keytabPath = keytabPath.some)
  }

  private def login(): Unit = {
    for (user <- principal; path <- keytabPath) {
      UserGroupInformation.loginUserFromKeytab(user, path.toString)
    }
  }

  /**
    * Returns a list of all files used by this hive source.
    *
    * @param includePartitionDirs if true then the partition directories will be included
    * @param includeTableDir      if true then the main table directory will be included
    * @return paths of all files and directories
    */
  def paths(includePartitionDirs: Boolean = false, includeTableDir: Boolean = false): List[Path] = {
    login()

    val files = ops.partitions(dbName, tableName).flatMap { partition =>
      val location = partition.getSd.getLocation
      val files = FilePattern(s"$location/*").toPaths()
      if (includePartitionDirs) {
        files :+ new Path(location)
      } else {
        files
      }
    }
    if (includeTableDir) {
      val location = spec().location
      files :+ new Path(location)
    } else {
      files
    }
  }

  def setPermissions(permission: FsPermission,
                     includePartitionDirs: Boolean = false,
                     includeTableDir: Boolean = false): Unit = {
    login()
    paths(includePartitionDirs, includeTableDir).foreach(fs.setPermission(_, permission))
  }

  /**
    * Sets the acl for all files of this hive source.
    * Even if the files are not located inside the table directory, this function will find them
    * and correctly update the spec.
    *
    * @param acl the acl values to set
    */
  def setAcl(acl: AclSpec,
             includePartitionDirs: Boolean = false,
             includeTableDir: Boolean = false): Unit = {
    login()
    paths(includePartitionDirs, includeTableDir).foreach { path =>
      HdfsSource(path).setAcl(acl)
    }
  }

  // returns the permission of the table location path
  def tablePermission(): FsPermission = {
    login()
    val location = ops.location(dbName, tableName)
    fs.getFileStatus(new Path(location)).getPermission
  }

  /**
    * Returns a TableSpec which contains details of the underlying table.
    * Similar to the Table class in the Hive API but using scala friendly types.
    */
  def spec(): TableSpec = {
    login()
    val table = client.getTable(dbName, tableName)
    val tableType = TableType.values().find(_.name.toLowerCase == table.getTableType.toLowerCase)
      .getOrElse(sys.error("Hive table type is not supported by this version of hive"))
    val params = table.getParameters.asScala.toMap
    TableSpec(
      tableName,
      tableType,
      table.getSd.getLocation,
      table.getSd.getNumBuckets,
      table.getSd.getBucketCols.asScala.toList,
      params,
      table.getSd.getInputFormat,
      table.getSd.getOutputFormat,
      table.getSd.getSerdeInfo.getName,
      table.getRetention,
      table.getCreateTime,
      table.getLastAccessTime,
      table.getOwner
    )
  }



  /**
   * The returned schema should take into account:
   *
   * 1) Any projection. If a projection is set, then it should return the schema in the same order
   * as the projection. If no projection is set then the schema should be driven from the hive metastore.
   *
   * 2) Any partitions set. These should be included in the schema columns.
   */
  override def schema(): Schema = {
    login()
    // if no field names were specified, then we will return the schema as is from the hive database,
    // otherwise we will keep only the requested fields
    val schema = if (projection.isEmpty) metastoreSchema
    else {
      // remember hive is always lower case, so when comparing requested field names with
      // hive fields we need to use lower case everything. And we need to return the schema
      // in the same order as the requested projection
      val columns = projection.map { fieldName =>
        val field = metastoreSchema.fields.find { it => it.name.equals(fieldName, true) }
        field.getOrElse(sys.error(s"Requested field $fieldName does not exist in the hive schema"))
      }
      Schema(columns)
    }

    schema
  }

  // returns the full underlying schema from the metastore including partition partitionKeys
  val metastoreSchema: Schema = {
    login()
    ops.schema(dbName, tableName)
  }

  //def  spec(): HiveSpec = HiveSpecFn.toHiveSpec(dbName, tableName)

  def isPartitionOnlyProjection(): Boolean = {
    val partitionKeyNames = HiveTable(dbName, tableName).partitionKeys().map(_.field.name)
    projection.nonEmpty && projection.map { it => it.toLowerCase() }.forall { it => partitionKeyNames.contains(it) }
  }

  override def parts(): List[Part] = {
    login()

    val table = client.getTable(dbName, tableName)
    val dialect = io.eels.component.hive.HiveDialect(table)
    val partitionKeys = HiveTable(dbName, tableName).partitionKeys()

    // all field names from the underlying hive schema
    val fieldNames = metastoreSchema.fields.map(_.name)

    // if we requested only partition columns, then we can get this information by scanning the file system
    // to see which partitions have been created. Then the presence of files in a partition location means that
    // that partition must have data.
    if (isPartitionOnlyProjection()) {
      logger.info("Requested projection contains only partitions; reading directly from metastore")
      // we pass in the schema so we can order the results to keep them aligned with the given projection
      List(new HivePartitionPart(dbName, tableName, fieldNames, Nil, metastoreSchema, predicate, dialect))
    } else {
      val files = HiveFilesFn(table, partitionKeys.map(_.field.name), constraints)
      logger.debug(s"Found ${files.size} visible hive files from all locations for $dbName:$tableName")

      // for each seperate hive file part we must pass in the metastore schema
      files.map { case (file, spec) =>
        new HiveFilePart(dialect, file, metastoreSchema, schema(), predicate, spec.parts.toList)
      }
    }
  }
}
{'fraction_non_alphanumeric': 0.08773573123516923, 'fraction_numerical': 0.0004371175221681029, 'mean_word_length': 3.158660088288756, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6613730', 'n_tokens_mistral': 4534, 'n_tokens_neox': 4430, 'n_words': 1439}