<commit_before>/*
 * The MIT License (MIT)
 * <p>
 * Copyright (c) 2018
 * <p>
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 * <p>
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 * <p>
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */
package io.techcode.streamy.elasticsearch.component.flow

import akka.NotUsed
import akka.actor.ActorSystem
import akka.stream._
import akka.stream.scaladsl.{Balance, Flow, GraphDSL, Merge, Source}
import akka.stream.stage._
import akka.util.ByteString
import com.softwaremill.sttp._
import io.techcode.streamy.elasticsearch.event.{ElasticsearchDropEvent, ElasticsearchFailureEvent, ElasticsearchPartialEvent, ElasticsearchSuccessEvent}
import io.techcode.streamy.util.json._

import scala.collection.mutable
import scala.concurrent.duration._
import scala.concurrent.{ExecutionContext, Future}
import scala.language.postfixOps
import scala.util.{Failure, Success}

/**
  * Elasticsearch flow companion.
  */
object ElasticsearchFlow {

  // Default values
  val DefaultBulk: Int = 500
  val DefaultWorker: Int = 1
  val DefaultRetry: FiniteDuration = 1 second

  // Component configuration
  case class Config(
    hosts: Seq[String],
    indexName: String,
    typeName: String,
    action: String,
    bulk: Int = DefaultBulk,
    worker: Int = DefaultWorker,
    retry: FiniteDuration = DefaultRetry
  )

  /**
    * Create a new elasticsearch flow.
    *
    * @param config sink configuration.
    * @return sink.
    */
  def apply(config: Config)(
    implicit httpClient: SttpBackend[Future, Source[ByteString, NotUsed]],
    system: ActorSystem,
    executionContext: ExecutionContext
  ): Flow[Json, Json, NotUsed] = {
    if (config.worker > 1) {
      balancer(Flow.fromGraph(new ElasticsearchFlowStage(config)), config.worker)
    } else {
      Flow.fromGraph(new ElasticsearchFlowStage(config))
    }
  }

  /**
    * Balancing jobs to a fixed pool of workers.
    *
    * @param worker      worker logic.
    * @param workerCount worker count.
    * @tparam In  input type.
    * @tparam Out output type.
    * @return balanced flow.
    */
  private def balancer[In, Out](worker: Flow[In, Out, Any], workerCount: Int): Flow[In, Out, NotUsed] = {
    import akka.stream.scaladsl.GraphDSL.Implicits._

    Flow.fromGraph(GraphDSL.create() { implicit b ⇒
      val balancer = b.add(Balance[In](workerCount, waitForAllDownstreams = true))
      val merge = b.add(Merge[Out](workerCount))

      for (_ ← 1 to workerCount) {
        // for each worker, add an edge from the balancer to the worker, then wire
        // it to the merge element
        balancer ~> worker.async ~> merge
      }

      FlowShape(balancer.in, merge.out)
    })
  }

  /**
    * Elasticsearch flow stage.
    *
    * @param config flow stage configuration.
    */
  private class ElasticsearchFlowStage(config: Config)(
    implicit val httpClient: SttpBackend[Future, Source[ByteString, NotUsed]],
    system: ActorSystem,
    executionContext: ExecutionContext
  ) extends GraphStage[FlowShape[Json, Json]] {

    // Inlet
    val in: Inlet[Json] = Inlet("ElasticsearchFlow.in")

    // Outlet
    val out: Outlet[Json] = Outlet("ElasticsearchFlow.out")

    // Shape
    override val shape: FlowShape[Json, Json] = FlowShape.of(in, out)

    // Logic generator
    override def createLogic(attr: Attributes): TimerGraphStageLogic = new ElasticsearchFlowLogic

    /**
      * Marshal all packets.
      *
      * @param pkts packets to process.
      * @return prepared request in bulk format.
      */
    private def marshalMessages(pkts: Seq[Json]): Array[Byte] = {
      pkts.map(marshalMessage)
        .reduce((x, y) => x ++ y)
        .compact.toArray[Byte]
    }

    /**
      * Marshal a single packet.
      *
      * @param pkt packet to marshal.
      * @return bytestring representation.
      */
    private def marshalMessage(pkt: Json): ByteString = {
      // Retrive header information
      val id = pkt.evaluate(Root / "_id").asString
      val `type` = pkt.evaluate(Root / "_type").asString.getOrElse(config.typeName)
      val version = pkt.evaluate(Root / "_version").asLong
      val versionType = pkt.evaluate(Root / "_version_type").asString

      // Build header
      val header = Json.obj(config.action -> {
        val builder = Json.objectBuilder()
          .put("_index" -> config.indexName)
          .put("_type" -> `type`)

        // Add version if present
        if (version.nonEmpty) {
          builder.put("_version" -> version.get)
        }

        // Add version type if present
        if (versionType.nonEmpty) {
          builder.put("_version_type" -> versionType.get)
        }

        // Add id if present
        if (id.nonEmpty) {
          builder.put("_id" -> id.get)
        }
        builder.result()
      })

      // Remove extra fields
      val doc = pkt.patch(Bulk(Root, Seq(
        Remove(Root / "_id", mustExist = false),
        Remove(Root / "_type", mustExist = false),
        Remove(Root / "_version", mustExist = false),
        Remove(Root / "_version_type", mustExist = false)
      ))).get
      println(header.toString)
      ByteString(header.toString()) ++ ByteString("\n") ++ ByteString(doc.toString()) ++ ByteString("\n")
    }
<commit_msg>[Removed] Useless debug<commit_after>/*
 * The MIT License (MIT)
 * <p>
 * Copyright (c) 2018
 * <p>
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 * <p>
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 * <p>
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */
package io.techcode.streamy.elasticsearch.component.flow

import akka.NotUsed
import akka.actor.ActorSystem
import akka.stream._
import akka.stream.scaladsl.{Balance, Flow, GraphDSL, Merge, Source}
import akka.stream.stage._
import akka.util.ByteString
import com.softwaremill.sttp._
import io.techcode.streamy.elasticsearch.event.{ElasticsearchDropEvent, ElasticsearchFailureEvent, ElasticsearchPartialEvent, ElasticsearchSuccessEvent}
import io.techcode.streamy.util.json._

import scala.collection.mutable
import scala.concurrent.duration._
import scala.concurrent.{ExecutionContext, Future}
import scala.language.postfixOps
import scala.util.{Failure, Success}

/**
  * Elasticsearch flow companion.
  */
object ElasticsearchFlow {

  // Default values
  val DefaultBulk: Int = 500
  val DefaultWorker: Int = 1
  val DefaultRetry: FiniteDuration = 1 second

  // Component configuration
  case class Config(
    hosts: Seq[String],
    indexName: String,
    typeName: String,
    action: String,
    bulk: Int = DefaultBulk,
    worker: Int = DefaultWorker,
    retry: FiniteDuration = DefaultRetry
  )

  /**
    * Create a new elasticsearch flow.
    *
    * @param config sink configuration.
    * @return sink.
    */
  def apply(config: Config)(
    implicit httpClient: SttpBackend[Future, Source[ByteString, NotUsed]],
    system: ActorSystem,
    executionContext: ExecutionContext
  ): Flow[Json, Json, NotUsed] = {
    if (config.worker > 1) {
      balancer(Flow.fromGraph(new ElasticsearchFlowStage(config)), config.worker)
    } else {
      Flow.fromGraph(new ElasticsearchFlowStage(config))
    }
  }

  /**
    * Balancing jobs to a fixed pool of workers.
    *
    * @param worker      worker logic.
    * @param workerCount worker count.
    * @tparam In  input type.
    * @tparam Out output type.
    * @return balanced flow.
    */
  private def balancer[In, Out](worker: Flow[In, Out, Any], workerCount: Int): Flow[In, Out, NotUsed] = {
    import akka.stream.scaladsl.GraphDSL.Implicits._

    Flow.fromGraph(GraphDSL.create() { implicit b ⇒
      val balancer = b.add(Balance[In](workerCount, waitForAllDownstreams = true))
      val merge = b.add(Merge[Out](workerCount))

      for (_ ← 1 to workerCount) {
        // for each worker, add an edge from the balancer to the worker, then wire
        // it to the merge element
        balancer ~> worker.async ~> merge
      }

      FlowShape(balancer.in, merge.out)
    })
  }

  /**
    * Elasticsearch flow stage.
    *
    * @param config flow stage configuration.
    */
  private class ElasticsearchFlowStage(config: Config)(
    implicit val httpClient: SttpBackend[Future, Source[ByteString, NotUsed]],
    system: ActorSystem,
    executionContext: ExecutionContext
  ) extends GraphStage[FlowShape[Json, Json]] {

    // Inlet
    val in: Inlet[Json] = Inlet("ElasticsearchFlow.in")

    // Outlet
    val out: Outlet[Json] = Outlet("ElasticsearchFlow.out")

    // Shape
    override val shape: FlowShape[Json, Json] = FlowShape.of(in, out)

    // Logic generator
    override def createLogic(attr: Attributes): TimerGraphStageLogic = new ElasticsearchFlowLogic

    /**
      * Marshal all packets.
      *
      * @param pkts packets to process.
      * @return prepared request in bulk format.
      */
    private def marshalMessages(pkts: Seq[Json]): Array[Byte] = {
      pkts.map(marshalMessage)
        .reduce((x, y) => x ++ y)
        .compact.toArray[Byte]
    }

    /**
      * Marshal a single packet.
      *
      * @param pkt packet to marshal.
      * @return bytestring representation.
      */
    private def marshalMessage(pkt: Json): ByteString = {
      // Retrive header information
      val id = pkt.evaluate(Root / "_id").asString
      val `type` = pkt.evaluate(Root / "_type").asString.getOrElse(config.typeName)
      val version = pkt.evaluate(Root / "_version").asLong
      val versionType = pkt.evaluate(Root / "_version_type").asString

      // Build header
      val header = Json.obj(config.action -> {
        val builder = Json.objectBuilder()
          .put("_index" -> config.indexName)
          .put("_type" -> `type`)

        // Add version if present
        if (version.nonEmpty) {
          builder.put("_version" -> version.get)
        }

        // Add version type if present
        if (versionType.nonEmpty) {
          builder.put("_version_type" -> versionType.get)
        }

        // Add id if present
        if (id.nonEmpty) {
          builder.put("_id" -> id.get)
        }
        builder.result()
      })

      // Remove extra fields
      val doc = pkt.patch(Bulk(Root, Seq(
        Remove(Root / "_id", mustExist = false),
        Remove(Root / "_type", mustExist = false),
        Remove(Root / "_version", mustExist = false),
        Remove(Root / "_version_type", mustExist = false)
      ))).get
      ByteString(header.toString()) ++ ByteString("\n") ++ ByteString(doc.toString()) ++ ByteString("\n")
    }

{'fraction_non_alphanumeric': 0.10588997812525318, 'fraction_numerical': 0.001782386777930811, 'mean_word_length': 2.9551425825056072, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 31, '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': '1571881', 'n_tokens_mistral': 3625, 'n_tokens_neox': 3568, 'n_words': 1105}