<commit_before>/*
 * Copyright 2017 Workday, Inc.
 *
 * This software is available under the MIT license.
 * Please see the LICENSE.txt file in this project.
 */

package com.workday.esclient

import com.fasterxml.jackson.annotation.{JsonIgnoreProperties, JsonProperty}
import com.fasterxml.jackson.databind.JsonNode
import com.google.common.annotations.VisibleForTesting
import com.google.gson.{JsonElement, JsonNull, JsonObject}
import io.searchbox.core._

import scala.collection.JavaConverters.{asJavaCollectionConverter, asScalaIteratorConverter}

/**
  * Trait wrapping Elasticsearch Query APIs
  */
trait EsQuery extends JestUtils {

  /**
    * Makes a basic query on Elasticsearch and returns an EsResult of the response.
    * @param index String ES index to query.
    * @param query String query content.
    * @return EsResult wrapping the response from ES.
    */
  def search(index: String, query: String): EsResult[EsSearchResponse] = {
    search(index, "", query)
  }

  /**
    * Makes a more sophisticated query on Elasticsearch including ES type names and parameters.
    * @param index String ES index to query.
    * @param typeName String ES type name.
    * @param query String query content.
    * @param params Map of params to include in ES request.
    * @return EsResult wrapping the response from ES.
    */
  def search(index: String, typeName: String = "", query: String, params: Map[String, Any] = Map()): EsResult[EsSearchResponse] = {
    val jestResult = jest.execute(buildSearchAction(query, index, typeName, params))
    handleSearchResult(jestResult)
  }

  /**
    * Makes a get request to Elasticsearch for the given id and index.
    * @param index String index to get from.
    * @param id String id of document to get.
    * @return EsResult wrapping the response from ES.
    */
  def get(index: String, id: String): EsResult[GetResponse] = {
    val jestResult = jest.execute(buildGetAction(index, id))
    handleJestResult(jestResult) { successfulJestResult =>
      JsonUtils.fromJson[GetResponse](successfulJestResult.getJsonString)
    }
  }

  /**
    * Makes a multi-get request to Elasticsearch for the given index and ids.
    * @param index String index to get from.
    * @param typeName String ES type name.
    * @param ids Sequence of ids of documents to get.
    * @param params Optional map of additional get parameters.
    * @return EsResult wrapping the response from ES.
    */
  def multiGet(index: String, typeName: String, ids: Seq[String], params: Map[String, Any] = Map()): EsResult[MultiGetResponse] = {
    // don't bother hitting elastic search if there are no id's
    if (ids.isEmpty){
      new EsResponse[MultiGetResponse](new MultiGetResponse(Nil))
    } else {
      val jestResult = jest.execute(buildMultiGetAction(index, typeName, ids.distinct, params))
      handleJestResult(jestResult, allowError = false) { successfulJestResult =>
        val parsedResult = JsonUtils.fromJson[MultiGetResponse](successfulJestResult.getJsonString)
        val docsIterator = successfulJestResult.getJsonObject.getAsJsonArray("docs").iterator.asScala
        def idAndSourcePair(jsonElement: JsonElement): (String, (Option[String], Option[String])) = {
          val jsonObject = jsonElement.getAsJsonObject
          val id = jsonObject.get(EsClient._ID)
          val source = Option(jsonObject.get(EsClient._SOURCE)).map(_.toString)
          val error = Option(jsonObject.get(EsClient.ERROR)).map(_.getAsString)
          id.getAsString ->(source, error)
        }
        val sourcesMappedById: Map[String, (Option[String], Option[String])] = docsIterator.map(idAndSourcePair).toMap
        def copyWithNewSource(getResponse: GetResponse): GetResponse = {
          sourcesMappedById(getResponse.id) match {
            case (source: Option[String], error: Option[String]) =>
              getResponse.copy(sourceIn = source, error = error)
          }
        }
        new MultiGetResponse(parsedResult.docs map copyWithNewSource)
      }
    }
  }

  /**
    * Makes a Multi Get request with the "_source" field provided as a parameter.
    * @param index String index to get from.
    * @param typeName String ES type name.
    * @param ids Sequence of ids of documents to get.
    * @param field String "_source" field to specify which documents to return source for.
    * @return EsResult wrapping the response from ES.
    */
  def multiGetSourceField(index: String, typeName: String, ids: Seq[String], field: String): EsResult[MultiGetResponse] = {
    multiGet(index, typeName, ids, Map("_source"->field))
  }

  /**
    * Gets the count of documents in a given index (or alias) with optional filtering by type.
    * @param index String index/alias name to get from.
    * @param typeName String optional ES type name (aka sid).
    * @return EsResult wrapping an integer count from ES.
    */
  def count(index: String, typeName: Option[String] = None): EsResult[Int] = {
    val indexBuilder = new Count.Builder().addIndex(index)
    val builder = typeName match {
      case Some(t) => indexBuilder.addType(t)
      case None => indexBuilder
    }
    val jestResult = jest.execute(builder.build())
    handleJestResult(jestResult) { _.getCount.toInt }
  }

  /**
    * Gets the count of search hits for a given query and index.
    * @param index String index to get from.
    * @param query String query content.
    * @return Int count of hits for given query.
    */
  def getSearchHitsCount(index: String, query: String): Int = {
    val jestResult = jest.execute(buildSearchAction(query, index))
    handleSearchResult(jestResult).get.hits.total
  }

  /**
    * Handles the Jest result from an ES request and returns a JSON object.
    * @param searchResult Jest SearchResult from ES.
    * @return EsResult wrapping the handled response from ES.
    */
  private[this] def handleSearchResult(searchResult: SearchResult): EsResult[EsSearchResponse] = {
    handleJestResult(searchResult) { successfulJestResult =>
      val json = successfulJestResult.getJsonObject

      EsSearchResponse(json.get("took").getAsLong, handleHitsInResult(json),
        if (json.has("aggregations")) json.get("aggregations").getAsJsonObject else new JsonObject
      )
    }
  }

  /**
    * Parses JSON search response and maps hits to SearchHit instances.
    * These instances have Index, Type, ID, Score, Source, and Matched fields.
    * @param searchJson JsonObject representing a search response from ES.
    * @return SearchHits object with hits data.
    */
  // scalastyle:off cyclomatic.complexity
  @VisibleForTesting
  private[esclient] def handleHitsInResult(searchJson: JsonObject): SearchHits = {
    val hitsObj = searchJson.get("hits").getAsJsonObject
    val hitsSeq = if (hitsObj.has("hits")) {
      hitsObj.get("hits").getAsJsonArray.iterator().asScala.toSeq.map(hit => {
        val h = hit.getAsJsonObject
        SearchHit(
          // index and typeName are sometimes left out for performance reasons with large result sets
          index = if (h.has(EsClient._INDEX)) h.get(EsClient._INDEX).getAsString else "",
          typeName = if (h.has(EsClient._TYPE)) h.get(EsClient._TYPE).getAsString else "",
          id = h.get(EsClient._ID).getAsString,
          score = if (h.has(EsClient._SCORE) && !h.get(EsClient._SCORE).isJsonNull)
            Some(h.get(EsClient._SCORE).getAsDouble) else None,
          source = if (h.has(EsClient._SOURCE)) h.get(EsClient._SOURCE).toString else "", // getting all doc IDs does not give back any fields
          matchedFields = {
            if (h.has(EsClient.MATCHED_QUERIES))
              Some(h.get(EsClient.MATCHED_QUERIES).getAsJsonArray.iterator().asScala.toSeq.map {
                element =>
                  // remove the named query delimiter and everything after it, if it exists
                  val queryName = element.getAsString
                  val delimIndex = queryName.indexOf(EsNames.NAME_QUERY_DELIMITER)
                  if (delimIndex >= 0) queryName.substring(0, delimIndex) else queryName
              })
            else
              None
          },
          explain = {
            if (h.has(EsClient._EXPLANATION)) {
              Some(h.get(EsClient._EXPLANATION).toString)<commit_msg>changing scope to protected for EsQuery methods<commit_after>/*
 * Copyright 2017 Workday, Inc.
 *
 * This software is available under the MIT license.
 * Please see the LICENSE.txt file in this project.
 */

package com.workday.esclient

import com.fasterxml.jackson.annotation.{JsonIgnoreProperties, JsonProperty}
import com.fasterxml.jackson.databind.JsonNode
import com.google.common.annotations.VisibleForTesting
import com.google.gson.{JsonElement, JsonNull, JsonObject}
import io.searchbox.core._

import scala.collection.JavaConverters.{asJavaCollectionConverter, asScalaIteratorConverter}

/**
  * Trait wrapping Elasticsearch Query APIs
  */
trait EsQuery extends JestUtils {

  /**
    * Makes a basic query on Elasticsearch and returns an EsResult of the response.
    * @param index String ES index to query.
    * @param query String query content.
    * @return EsResult wrapping the response from ES.
    */
  def search(index: String, query: String): EsResult[EsSearchResponse] = {
    search(index, "", query)
  }

  /**
    * Makes a more sophisticated query on Elasticsearch including ES type names and parameters.
    * @param index String ES index to query.
    * @param typeName String ES type name.
    * @param query String query content.
    * @param params Map of params to include in ES request.
    * @return EsResult wrapping the response from ES.
    */
  def search(index: String, typeName: String = "", query: String, params: Map[String, Any] = Map()): EsResult[EsSearchResponse] = {
    val jestResult = jest.execute(buildSearchAction(query, index, typeName, params))
    handleSearchResult(jestResult)
  }

  /**
    * Makes a get request to Elasticsearch for the given id and index.
    * @param index String index to get from.
    * @param id String id of document to get.
    * @return EsResult wrapping the response from ES.
    */
  def get(index: String, id: String): EsResult[GetResponse] = {
    val jestResult = jest.execute(buildGetAction(index, id))
    handleJestResult(jestResult) { successfulJestResult =>
      JsonUtils.fromJson[GetResponse](successfulJestResult.getJsonString)
    }
  }

  /**
    * Makes a multi-get request to Elasticsearch for the given index and ids.
    * @param index String index to get from.
    * @param typeName String ES type name.
    * @param ids Sequence of ids of documents to get.
    * @param params Optional map of additional get parameters.
    * @return EsResult wrapping the response from ES.
    */
  def multiGet(index: String, typeName: String, ids: Seq[String], params: Map[String, Any] = Map()): EsResult[MultiGetResponse] = {
    // don't bother hitting elastic search if there are no id's
    if (ids.isEmpty){
      new EsResponse[MultiGetResponse](new MultiGetResponse(Nil))
    } else {
      val jestResult = jest.execute(buildMultiGetAction(index, typeName, ids.distinct, params))
      handleJestResult(jestResult, allowError = false) { successfulJestResult =>
        val parsedResult = JsonUtils.fromJson[MultiGetResponse](successfulJestResult.getJsonString)
        val docsIterator = successfulJestResult.getJsonObject.getAsJsonArray("docs").iterator.asScala
        def idAndSourcePair(jsonElement: JsonElement): (String, (Option[String], Option[String])) = {
          val jsonObject = jsonElement.getAsJsonObject
          val id = jsonObject.get(EsClient._ID)
          val source = Option(jsonObject.get(EsClient._SOURCE)).map(_.toString)
          val error = Option(jsonObject.get(EsClient.ERROR)).map(_.getAsString)
          id.getAsString ->(source, error)
        }
        val sourcesMappedById: Map[String, (Option[String], Option[String])] = docsIterator.map(idAndSourcePair).toMap
        def copyWithNewSource(getResponse: GetResponse): GetResponse = {
          sourcesMappedById(getResponse.id) match {
            case (source: Option[String], error: Option[String]) =>
              getResponse.copy(sourceIn = source, error = error)
          }
        }
        new MultiGetResponse(parsedResult.docs map copyWithNewSource)
      }
    }
  }

  /**
    * Makes a Multi Get request with the "_source" field provided as a parameter.
    * @param index String index to get from.
    * @param typeName String ES type name.
    * @param ids Sequence of ids of documents to get.
    * @param field String "_source" field to specify which documents to return source for.
    * @return EsResult wrapping the response from ES.
    */
  def multiGetSourceField(index: String, typeName: String, ids: Seq[String], field: String): EsResult[MultiGetResponse] = {
    multiGet(index, typeName, ids, Map("_source"->field))
  }

  /**
    * Gets the count of documents in a given index (or alias) with optional filtering by type.
    * @param index String index/alias name to get from.
    * @param typeName String optional ES type name (aka sid).
    * @return EsResult wrapping an integer count from ES.
    */
  def count(index: String, typeName: Option[String] = None): EsResult[Int] = {
    val indexBuilder = new Count.Builder().addIndex(index)
    val builder = typeName match {
      case Some(t) => indexBuilder.addType(t)
      case None => indexBuilder
    }
    val jestResult = jest.execute(builder.build())
    handleJestResult(jestResult) { _.getCount.toInt }
  }

  /**
    * Gets the count of search hits for a given query and index.
    * @param index String index to get from.
    * @param query String query content.
    * @return Int count of hits for given query.
    */
  def getSearchHitsCount(index: String, query: String): Int = {
    val jestResult = jest.execute(buildSearchAction(query, index))
    handleSearchResult(jestResult).get.hits.total
  }

  /**
    * Handles the Jest result from an ES request and returns a JSON object.
    * @param searchResult Jest SearchResult from ES.
    * @return EsResult wrapping the handled response from ES.
    */
  protected def handleSearchResult(searchResult: SearchResult): EsResult[EsSearchResponse] = {
    handleJestResult(searchResult) { successfulJestResult =>
      val json = successfulJestResult.getJsonObject

      EsSearchResponse(json.get("took").getAsLong, handleHitsInResult(json),
        if (json.has("aggregations")) json.get("aggregations").getAsJsonObject else new JsonObject
      )
    }
  }

  /**
    * Parses JSON search response and maps hits to SearchHit instances.
    * These instances have Index, Type, ID, Score, Source, and Matched fields.
    * @param searchJson JsonObject representing a search response from ES.
    * @return SearchHits object with hits data.
    */
  // scalastyle:off cyclomatic.complexity
  @VisibleForTesting
  protected[esclient] def handleHitsInResult(searchJson: JsonObject): SearchHits = {
    val hitsObj = searchJson.get("hits").getAsJsonObject
    val hitsSeq = if (hitsObj.has("hits")) {
      hitsObj.get("hits").getAsJsonArray.iterator().asScala.toSeq.map(hit => {
        val h = hit.getAsJsonObject
        SearchHit(
          // index and typeName are sometimes left out for performance reasons with large result sets
          index = if (h.has(EsClient._INDEX)) h.get(EsClient._INDEX).getAsString else "",
          typeName = if (h.has(EsClient._TYPE)) h.get(EsClient._TYPE).getAsString else "",
          id = h.get(EsClient._ID).getAsString,
          score = if (h.has(EsClient._SCORE) && !h.get(EsClient._SCORE).isJsonNull)
            Some(h.get(EsClient._SCORE).getAsDouble) else None,
          source = if (h.has(EsClient._SOURCE)) h.get(EsClient._SOURCE).toString else "", // getting all doc IDs does not give back any fields
          matchedFields = {
            if (h.has(EsClient.MATCHED_QUERIES))
              Some(h.get(EsClient.MATCHED_QUERIES).getAsJsonArray.iterator().asScala.toSeq.map {
                element =>
                  // remove the named query delimiter and everything after it, if it exists
                  val queryName = element.getAsString
                  val delimIndex = queryName.indexOf(EsNames.NAME_QUERY_DELIMITER)
                  if (delimIndex >= 0) queryName.substring(0, delimIndex) else queryName
              })
            else
              None
          },
          explain = {
            if (h.has(EsClient._EXPLANATION)) {
              Some(h.get(EsClient._EXPLANATION).toString)
{'fraction_non_alphanumeric': 0.09785231660231661, 'fraction_numerical': 0.0007239382239382239, 'mean_word_length': 3.1640291384074355, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 25, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 4}, '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': '4933090', 'n_tokens_mistral': 4719, 'n_tokens_neox': 4609, 'n_words': 1415}