<commit_before>package io.flow.event


import io.flow.play.util.{FlowEnvironment, Random}
import com.amazonaws.auth.BasicAWSCredentials
import com.amazonaws.services.kinesis.AmazonKinesisClient
import com.amazonaws.services.kinesis.model._
import play.api.libs.json.{JsValue, Json}
import play.api.Logger

import scala.reflect.runtime.universe._
import scala.concurrent.{ExecutionContext, Future}
import scala.util.{Failure, Success, Try}
import java.nio.ByteBuffer

import collection.JavaConverters._

trait Queue {
  def stream[T: TypeTag](implicit ec: ExecutionContext): Stream
}

trait Stream {
  def publish(event: JsValue)
  def consume(f: JsValue => Unit)(implicit ec: ExecutionContext)
}

class KinesisQueue @javax.inject.Inject() (
  config: io.flow.play.util.Config
) extends Queue {

  private[this] val credentials = new BasicAWSCredentials(
    config.requiredString("aws.access.key"),
    config.requiredString("aws.secret.key")
  )

  // Get shard count from config - default to 1 if not available, for backward compatibility with existing streams
  private[this] val numberShards = config.optionalPositiveInt("aws.kinesis.shard.count").getOrElse(1)

  private[this] val client = new AmazonKinesisClient(credentials)
  var kinesisStreams: scala.collection.mutable.Map[String, KinesisStream] = scala.collection.mutable.Map[String, KinesisStream]()

  override def stream[T: TypeTag](implicit ec: ExecutionContext): Stream = {
    val name = typeOf[T].toString

    StreamNames(FlowEnvironment.Current).json(name) match {
      case None => {
        name match {
          case "Any" => {
            sys.error(s"FlowKinesisError Stream[$name] In order to consume events, you must annotate the type you are expecting as this is used to build the stream. Type should be something like io.flow.user.v0.unions.SomeEvent")
          }
          case _ => {
            sys.error(s"FlowKinesisError Stream[$name] Could not parse stream name. Expected something like io.flow.user.v0.unions.SomeEvent")
          }
        }
      }
      case Some(streamName) => {
        if (!kinesisStreams.contains(streamName))
          kinesisStreams.put(streamName, KinesisStream(client, streamName, numberShards))

        kinesisStreams.get(streamName).getOrElse(KinesisStream(client, streamName, numberShards))
      }
    }
  }

}

case class KinesisStream(
  client: AmazonKinesisClient,
  name: String,
  numberShards: Int = 1
) (
  implicit ec: ExecutionContext
) extends Stream {

  private[this] val recordLimit = 1000
  private[this] var shardSequenceNumberMap = scala.collection.mutable.Map.empty[String,String]

  setup

  def publish(event: JsValue) {
    val partitionKey = Random().alphaNumeric(30)
    val data = Json.stringify(event)
    insertMessage(partitionKey, data)
  }

  def consume(f: JsValue => Unit)(implicit ec: ExecutionContext) {
    processMessages(f).recover {
      case e: Throwable => {
        Logger.error(s"FlowKinesisError Stream[$name] Error processing: $e")
      }
    }
  }

  /**
   * Publish Helper Functions
   **/
  def insertMessage(partitionKey: String, data: String) {
    try {
      client.putRecord(
        new PutRecordRequest()
          .withData(ByteBuffer.wrap(data.getBytes))
          .withPartitionKey(partitionKey)
          .withStreamName(name)<commit_msg>ensure UTF-8 character encoding on published message (#106)<commit_after>package io.flow.event


import io.flow.play.util.{FlowEnvironment, Random}
import com.amazonaws.auth.BasicAWSCredentials
import com.amazonaws.services.kinesis.AmazonKinesisClient
import com.amazonaws.services.kinesis.model._
import play.api.libs.json.{JsValue, Json}
import play.api.Logger

import scala.reflect.runtime.universe._
import scala.concurrent.{ExecutionContext, Future}
import scala.util.{Failure, Success, Try}
import java.nio.ByteBuffer

import collection.JavaConverters._

trait Queue {
  def stream[T: TypeTag](implicit ec: ExecutionContext): Stream
}

trait Stream {
  def publish(event: JsValue)
  def consume(f: JsValue => Unit)(implicit ec: ExecutionContext)
}

class KinesisQueue @javax.inject.Inject() (
  config: io.flow.play.util.Config
) extends Queue {

  private[this] val credentials = new BasicAWSCredentials(
    config.requiredString("aws.access.key"),
    config.requiredString("aws.secret.key")
  )

  // Get shard count from config - default to 1 if not available, for backward compatibility with existing streams
  private[this] val numberShards = config.optionalPositiveInt("aws.kinesis.shard.count").getOrElse(1)

  private[this] val client = new AmazonKinesisClient(credentials)
  var kinesisStreams: scala.collection.mutable.Map[String, KinesisStream] = scala.collection.mutable.Map[String, KinesisStream]()

  override def stream[T: TypeTag](implicit ec: ExecutionContext): Stream = {
    val name = typeOf[T].toString

    StreamNames(FlowEnvironment.Current).json(name) match {
      case None => {
        name match {
          case "Any" => {
            sys.error(s"FlowKinesisError Stream[$name] In order to consume events, you must annotate the type you are expecting as this is used to build the stream. Type should be something like io.flow.user.v0.unions.SomeEvent")
          }
          case _ => {
            sys.error(s"FlowKinesisError Stream[$name] Could not parse stream name. Expected something like io.flow.user.v0.unions.SomeEvent")
          }
        }
      }
      case Some(streamName) => {
        if (!kinesisStreams.contains(streamName))
          kinesisStreams.put(streamName, KinesisStream(client, streamName, numberShards))

        kinesisStreams.get(streamName).getOrElse(KinesisStream(client, streamName, numberShards))
      }
    }
  }

}

case class KinesisStream(
  client: AmazonKinesisClient,
  name: String,
  numberShards: Int = 1
) (
  implicit ec: ExecutionContext
) extends Stream {

  private[this] val recordLimit = 1000
  private[this] var shardSequenceNumberMap = scala.collection.mutable.Map.empty[String,String]

  setup

  def publish(event: JsValue) {
    val partitionKey = Random().alphaNumeric(30)
    val data = Json.stringify(event)
    insertMessage(partitionKey, data)
  }

  def consume(f: JsValue => Unit)(implicit ec: ExecutionContext) {
    processMessages(f).recover {
      case e: Throwable => {
        Logger.error(s"FlowKinesisError Stream[$name] Error processing: $e")
      }
    }
  }

  /**
   * Publish Helper Functions
   **/
  def insertMessage(partitionKey: String, data: String) {
    try {
      client.putRecord(
        new PutRecordRequest()
          .withData(ByteBuffer.wrap(data.getBytes("UTF-8")))
          .withPartitionKey(partitionKey)
          .withStreamName(name)
{'fraction_non_alphanumeric': 0.1053024645257655, 'fraction_numerical': 0.004032860343539955, 'mean_word_length': 3.7355021216407356, 'pattern_counts': {'":': 0, '<': 3, '<?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': '1854626', 'n_tokens_mistral': 2075, 'n_tokens_neox': 2042, 'n_words': 464}