<commit_before>package com.workday.esclient

import com.google.common.annotations.VisibleForTesting
import com.google.gson.JsonObject
import io.searchbox.client.JestResult
import io.searchbox.core.{Search, SearchScroll}
import io.searchbox.params.Parameters

/**
  * Trait wrapping Elasticsearch Scan and Scroll APIs.
  */
trait EsScanAndScroll extends JestUtils with EsQuery {
  private[this] val emptyScanAndScrollResponse: EsResult[ScanAndScrollResponse] =
    EsResponse(ScanAndScrollResponse(0, SearchHits(0, None, Seq.empty), new JsonObject, ""))
  private[this] val UnlimitedResultsDefaultParams = Map(EsNames.FilterPath -> EsNames.FilterPathParams)

  /**
    * Performs a ScanAndScroll search, combines all paged results into an iterator and
    * aggregates each result in the iterator into a single EsResult[EsSearchResponse].
    * @param index String index to search
    * @param typeName String index type
    * @param query String query to perform
    * @param scanAndScrollSize Int size of each scan and scroll response
    * @return EsResult[EsSearchResponse] of aggregate result from scan-and-scrolled results
    */
  def unlimitedSearch(index: String, typeName: String = "", query: String, scanAndScrollSize: Int): EsResult[EsSearchResponse] = {
    val unlimitedResultsParams = UnlimitedResultsDefaultParams ++ Map(Parameters.SIZE -> scanAndScrollSize)
    val scanAndScrollSearch = createScrolledSearch(index, typeName, query, unlimitedResultsParams)
    val scanAndScrollIter = getScrolledSearchIterator(scanAndScrollSearch, true)
    val scanAndScrollResult = scanAndScrollIter.foldLeft(emptyScanAndScrollResponse)(accumulateScanAndScrollResponses)

    // Turn ScanAndScroll result into an EsResult
    scanAndScrollResult match {
      case EsResponse(response) => EsResponse(EsSearchResponse(response.took, response.hits, response.aggregations))
      case e: EsError => e
      case e: EsInvalidResponse => e
    }
  }

  /**
    * Does a scan and scroll search and returns an iterator over the results.
    * (https://www.elastic.co/guide/en/elasticsearch/guide/current/scan-scroll.html)
    * @param index String index to search
    * @param typeName String index type
    * @param query String query to perform
    * @param params Map[String,Any] query parameters to include; defaults to empty Map
    * @return Iterator over EsResults of SearchHits
    */
  def createScrolledSearchIterator(index: String, typeName: String, query: String, params: Map[String, Any] = Map()): Iterator[EsResult[SearchHits]] = {
    val result = createScrolledSearch(index, typeName, query, params)
    getScrolledSearchHitsIterator(result)
  }

  def createScrolledSearch(index: String, typeName: String = "", query: String, params: Map[String, Any] = Map()): EsResult[ScanAndScrollResponse] = {
    val jestResult : JestResult = jest.execute(buildScrolledSearchAction(query, index, typeName, params))
    handleScrollResult(jestResult)
  }

  def getScrolledSearchResults(scrollID: String): EsResult[ScanAndScrollResponse] = {
    val jestResult : JestResult = jest.execute(buildGetScrolledSearchResultsAction(scrollID))
    handleScrollResult(jestResult)
  }

  private[this] def handleScrollResult(scrollResult: JestResult): EsResult[ScanAndScrollResponse] = {
    handleJestResult(scrollResult) { successfulJestResult =>
      val json = successfulJestResult.getJsonObject
      val aggregations = Option(json.get("aggregations")).getOrElse(new JsonObject)
      ScanAndScrollResponse(json.get("took").getAsLong, handleHitsInResult(json), aggregations.getAsJsonObject, json.get(EsClient._SCROLL_ID).getAsString())
    }
  }

  /**
    * Adds the values of the nextScanAndScrollResult onto the accumulator.
    * If we encounter an error in the scan and scroll search, just return the error and stop accumulating.
    * @param scanAndScrollResultAccumulator tuple of current accumulated results and boolean of whether an error has been encountered already
    * @param nextScanAndScrollResult the next "page" in the ScanAndScroll search result
    * @return the current accumulated response and whether we should continue accumulating
    */
  @VisibleForTesting
  private[workday] def accumulateScanAndScrollResponses(scanAndScrollResultAccumulator: EsResult[ScanAndScrollResponse],
    nextScanAndScrollResult: EsResult[ScanAndScrollResponse]): EsResult[ScanAndScrollResponse] = {

    (scanAndScrollResultAccumulator, nextScanAndScrollResult) match {
      case (EsResponse(response1), EsResponse(response2)) =>
        EsResponse(ScanAndScrollResponse(
          response1.took + response2.took,
          addSearchHits(response1.hits, response2.hits),
          // Take the first non-null aggregation of the responses. Theoretically, only the first page will have an aggregation:
          // https://www.elastic.co/guide/en/elasticsearch/reference/current/search-request-scroll.html
          if (response1.aggregations.entrySet().isEmpty) response2.aggregations else response1.aggregations,
          response1.scrollID
        ))
      case (EsResponse(response), e) => e<commit_msg>finishing scan and scroll docs<commit_after>package com.workday.esclient

import com.google.common.annotations.VisibleForTesting
import com.google.gson.JsonObject
import io.searchbox.client.JestResult
import io.searchbox.core.{Search, SearchScroll}
import io.searchbox.params.Parameters

/**
  * Trait wrapping Elasticsearch Scan and Scroll APIs.
  */
trait EsScanAndScroll extends JestUtils with EsQuery {
  private[this] val emptyScanAndScrollResponse: EsResult[ScanAndScrollResponse] =
    EsResponse(ScanAndScrollResponse(0, SearchHits(0, None, Seq.empty), new JsonObject, ""))
  private[this] val UnlimitedResultsDefaultParams = Map(EsNames.FilterPath -> EsNames.FilterPathParams)

  /**
    * Performs a ScanAndScroll search, combines all paged results into an iterator and
    * aggregates each result in the iterator into a single EsResult[EsSearchResponse].
    * @param index String index to search
    * @param typeName String index type
    * @param query String query to perform
    * @param scanAndScrollSize Int size of each scan and scroll response
    * @return EsResult[EsSearchResponse] of aggregate result from scan-and-scrolled results
    */
  def unlimitedSearch(index: String, typeName: String = "", query: String, scanAndScrollSize: Int): EsResult[EsSearchResponse] = {
    val unlimitedResultsParams = UnlimitedResultsDefaultParams ++ Map(Parameters.SIZE -> scanAndScrollSize)
    val scanAndScrollSearch = createScrolledSearch(index, typeName, query, unlimitedResultsParams)
    val scanAndScrollIter = getScrolledSearchIterator(scanAndScrollSearch, true)
    val scanAndScrollResult = scanAndScrollIter.foldLeft(emptyScanAndScrollResponse)(accumulateScanAndScrollResponses)

    // Turn ScanAndScroll result into an EsResult
    scanAndScrollResult match {
      case EsResponse(response) => EsResponse(EsSearchResponse(response.took, response.hits, response.aggregations))
      case e: EsError => e
      case e: EsInvalidResponse => e
    }
  }

  /**
    * Does a scan and scroll search and returns an iterator over the results.
    * (https://www.elastic.co/guide/en/elasticsearch/guide/current/scan-scroll.html)
    * @param index String index to search
    * @param typeName String index type
    * @param query String query to perform
    * @param params Map[String,Any] query parameters to include; defaults to empty Map
    * @return Iterator over EsResults of SearchHits
    */
  def createScrolledSearchIterator(index: String, typeName: String, query: String, params: Map[String, Any] = Map()): Iterator[EsResult[SearchHits]] = {
    val result = createScrolledSearch(index, typeName, query, params)
    getScrolledSearchHitsIterator(result)
  }

  /**
    * Creates a scrolled search and returns the Jest response.
    * @param index String index to search on
    * @param typeName String type name to search
    * @param query String query to search on
    * @param params String params for search query
    * @return EsResult of Elasticsearch scan and scroll response
    */
  def createScrolledSearch(index: String, typeName: String = "", query: String, params: Map[String, Any] = Map()): EsResult[ScanAndScrollResponse] = {
    val jestResult : JestResult = jest.execute(buildScrolledSearchAction(query, index, typeName, params))
    handleScrollResult(jestResult)
  }

  /**
    * Gets the scrolled search result using a scroll ID
    * @param scrollID String provided scroll ID
    * @return EsResult of Elasticsearch scan and scroll response
    */
  def getScrolledSearchResults(scrollID: String): EsResult[ScanAndScrollResponse] = {
    val jestResult : JestResult = jest.execute(buildGetScrolledSearchResultsAction(scrollID))
    handleScrollResult(jestResult)
  }

  /**
    * Returns the Elasticsearch scan and scroll response from the Jest client
    * @param scrollResult Jest result to be handled
    * @return EsResult of Elasticsearch scan and scroll response
    */
  private[this] def handleScrollResult(scrollResult: JestResult): EsResult[ScanAndScrollResponse] = {
    handleJestResult(scrollResult) { successfulJestResult =>
      val json = successfulJestResult.getJsonObject
      val aggregations = Option(json.get("aggregations")).getOrElse(new JsonObject)
      ScanAndScrollResponse(json.get("took").getAsLong, handleHitsInResult(json), aggregations.getAsJsonObject, json.get(EsClient._SCROLL_ID).getAsString())
    }
  }

  /**
    * Adds the values of the nextScanAndScrollResult onto the accumulator.
    * If we encounter an error in the scan and scroll search, just return the error and stop accumulating.
    * @param scanAndScrollResultAccumulator tuple of current accumulated results and boolean of whether an error has been encountered already
    * @param nextScanAndScrollResult the next "page" in the ScanAndScroll search result
    * @return the current accumulated response and whether we should continue accumulating
    */
  @VisibleForTesting
  private[workday] def accumulateScanAndScrollResponses(scanAndScrollResultAccumulator: EsResult[ScanAndScrollResponse],
    nextScanAndScrollResult: EsResult[ScanAndScrollResponse]): EsResult[ScanAndScrollResponse] = {

    (scanAndScrollResultAccumulator, nextScanAndScrollResult) match {
      case (EsResponse(response1), EsResponse(response2)) =>
        EsResponse(ScanAndScrollResponse(
          response1.took + response2.took,
          addSearchHits(response1.hits, response2.hits),
          // Take the first non-null aggregation of the responses. Theoretically, only the first page will have an aggregation:
          // https://www.elastic.co/guide/en/elasticsearch/reference/current/search-request-scroll.html
          if (response1.aggregations.entrySet().isEmpty) response2.aggregations else response1.aggregations,
          response1.scrollID
        ))
      case (EsResponse(response), e) => e
{'fraction_non_alphanumeric': 0.08117238860482104, 'fraction_numerical': 0.002191380569758948, 'mean_word_length': 4.58256880733945, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 19, 'https://': 4, 'lorem ipsum': 0, 'www.': 4, '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': '5324891', 'n_tokens_mistral': 2984, 'n_tokens_neox': 2900, 'n_words': 843}