<commit_before>package pl.touk.nussknacker.engine.kafka

import org.apache.flink.api.common.typeinfo.TypeInformation
import org.apache.flink.streaming.api.functions.TimestampAssigner
import org.apache.flink.streaming.api.functions.source.SourceFunction
import org.apache.flink.streaming.connectors.kafka.FlinkKafkaConsumer09
import org.apache.flink.streaming.util.serialization.DeserializationSchema
import pl.touk.nussknacker.engine.api.process.{Source, TestDataGenerator}
import pl.touk.nussknacker.engine.api.test.{TestDataParser, TestDataSplit}
import pl.touk.nussknacker.engine.api.{MetaData, MethodToInvoke, ParamName}
import pl.touk.nussknacker.engine.flink.api.process.{FlinkSource, FlinkSourceFactory}
import pl.touk.nussknacker.engine.kafka.KafkaSourceFactory._

/**<pre>
  * Wrapper for [[org.apache.flink.streaming.connectors.kafka.FlinkKafkaConsumer09]]
  * Features:
  *   - fetch latest N records which can be later used to test process in UI
  *     Fetching data is defined in [[pl.touk.nussknacker.engine.kafka.KafkaSourceFactory.KafkaSource]] which
  *     extends [[pl.touk.nussknacker.engine.api.process.TestDataGenerator]]. See [[pl.touk.nussknacker.engine.kafka.KafkaEspUtils#readLastMessages]]
  *   - reset Kafka's offset to latest value - `forceLatestRead` property, see [[pl.touk.nussknacker.engine.kafka.KafkaEspUtils#setOffsetToLatest]]
  *</pre>
* */
class KafkaSourceFactory[T: TypeInformation](config: KafkaConfig,
                                             schema: DeserializationSchema[T],
                                             val timestampAssigner: Option[TimestampAssigner[T]],
                                             testPrepareInfo: TestDataSplit) extends FlinkSourceFactory[T] with Serializable {

  @MethodToInvoke
  def create(processMetaData: MetaData, @ParamName(`TopicParamName`) topic: String): Source[T] with TestDataGenerator = {
    val espKafkaProperties = config.kafkaEspProperties.getOrElse(Map.empty)
    if (espKafkaProperties.get("forceLatestRead").exists(java.lang.Boolean.parseBoolean)) {
      KafkaEspUtils.setOffsetToLatest(topic, processMetaData.id, config)
    } else {
      ()
    }
    new KafkaSource(consumerGroupId = processMetaData.id, topic = topic)
  }

  override def testDataParser: Option[TestDataParser[T]] = Some(new TestDataParser[T] {
    override def parseTestData(data: Array[Byte]) = testPrepareInfo.splitData(data).map(schema.deserialize)
  })

  class KafkaSource(consumerGroupId: String, topic: String) extends FlinkSource[T] with Serializable
    with TestDataGenerator {
    override def typeInformation: TypeInformation[T] =
      implicitly[TypeInformation[T]]

    override def toFlinkSource: SourceFunction[T] = {
      new FlinkKafkaConsumer09[T](topic, schema, KafkaEspUtils.toProperties(config, Some(consumerGroupId)))
    }

    override def generateTestData(size: Int) =
      testPrepareInfo.joinData(KafkaEspUtils.readLastMessages(topic, size, config))

    override def timestampAssigner = KafkaSourceFactory.this.timestampAssigner
  }

}

object KafkaSourceFactory {

  final val TopicParamName = "topic"

}
<commit_msg>KafkaSourceFactory can have topic preconfigured<commit_after>package pl.touk.nussknacker.engine.kafka

import org.apache.flink.api.common.typeinfo.TypeInformation
import org.apache.flink.streaming.api.functions.TimestampAssigner
import org.apache.flink.streaming.api.functions.source.SourceFunction
import org.apache.flink.streaming.connectors.kafka.FlinkKafkaConsumer09
import org.apache.flink.streaming.util.serialization.DeserializationSchema
import pl.touk.nussknacker.engine.api.process.{Source, TestDataGenerator}
import pl.touk.nussknacker.engine.api.test.{TestDataParser, TestDataSplit}
import pl.touk.nussknacker.engine.api.{MetaData, MethodToInvoke, ParamName}
import pl.touk.nussknacker.engine.flink.api.process.{FlinkSource, FlinkSourceFactory}
import pl.touk.nussknacker.engine.kafka.KafkaSourceFactory._

/** <pre>
  * Wrapper for [[org.apache.flink.streaming.connectors.kafka.FlinkKafkaConsumer09]]
  * Features:
  *   - fetch latest N records which can be later used to test process in UI
  * Fetching data is defined in [[pl.touk.nussknacker.engine.kafka.BaseKafkaSourceFactory.KafkaSource]] which
  * extends [[pl.touk.nussknacker.engine.api.process.TestDataGenerator]]. See [[pl.touk.nussknacker.engine.kafka.KafkaEspUtils#readLastMessages]]
  *   - reset Kafka's offset to latest value - `forceLatestRead` property, see [[pl.touk.nussknacker.engine.kafka.KafkaEspUtils#setOffsetToLatest]]
  *
  * BaseKafkaSourceFactory comes in two variants:
  *   - KafkaSourceFactory - `topic` parameter has to be passed on frontend
  *   - SingleTopicKafkaSourceFactory - topic is defined on level of configuration
  *
  * </pre>
  * */
class KafkaSourceFactory[T: TypeInformation](config: KafkaConfig,
                                             schema: DeserializationSchema[T],
                                             timestampAssigner: Option[TimestampAssigner[T]],
                                             testPrepareInfo: TestDataSplit) extends BaseKafkaSourceFactory(config, schema, timestampAssigner, testPrepareInfo) {

  @MethodToInvoke
  def create(processMetaData: MetaData, @ParamName(`TopicParamName`) topic: String): Source[T] with TestDataGenerator = {
    createSource(processMetaData, topic)
  }

}

object KafkaSourceFactory {

  final val TopicParamName = "topic"

}


class SingleTopicKafkaSourceFactory[T: TypeInformation](topic: String,
                                                        config: KafkaConfig,
                                                        schema: DeserializationSchema[T],
                                                        timestampAssigner: Option[TimestampAssigner[T]],
                                                        testPrepareInfo: TestDataSplit) extends BaseKafkaSourceFactory(config, schema, timestampAssigner, testPrepareInfo) {

  @MethodToInvoke
  def create(processMetaData: MetaData): Source[T] with TestDataGenerator = {
    createSource(processMetaData, topic)
  }

}

abstract class BaseKafkaSourceFactory[T: TypeInformation](config: KafkaConfig,
                                                          schema: DeserializationSchema[T],
                                                          val timestampAssigner: Option[TimestampAssigner[T]],
                                                          testPrepareInfo: TestDataSplit) extends FlinkSourceFactory[T] with Serializable {


  override def testDataParser: Option[TestDataParser[T]] = Some(new TestDataParser[T] {
    override def parseTestData(data: Array[Byte]): List[T] = testPrepareInfo.splitData(data).map(schema.deserialize)
  })

  protected def createSource(processMetaData: MetaData, topic: String): KafkaSource = {
    val espKafkaProperties = config.kafkaEspProperties.getOrElse(Map.empty)
    if (espKafkaProperties.get("forceLatestRead").exists(java.lang.Boolean.parseBoolean)) {
      KafkaEspUtils.setOffsetToLatest(topic, processMetaData.id, config)
    } else {
      ()
    }
    new KafkaSource(consumerGroupId = processMetaData.id, topic = topic)
  }

  class KafkaSource(consumerGroupId: String, topic: String) extends FlinkSource[T] with Serializable with TestDataGenerator {

    override def typeInformation: TypeInformation[T] =
      implicitly[TypeInformation[T]]

    override def toFlinkSource: SourceFunction[T] = {
      new FlinkKafkaConsumer09[T](topic, schema, KafkaEspUtils.toProperties(config, Some(consumerGroupId)))
    }

    override def generateTestData(size: Int): Array[Byte] =
      testPrepareInfo.joinData(KafkaEspUtils.readLastMessages(topic, size, config))

    override def timestampAssigner: Option[TimestampAssigner[T]] = BaseKafkaSourceFactory.this.timestampAssigner
  }

}
{'fraction_non_alphanumeric': 0.09304703476482618, 'fraction_numerical': 0.0015337423312883436, 'mean_word_length': 3.592136150234742, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 26, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3026124', 'n_tokens_mistral': 2259, 'n_tokens_neox': 2212, 'n_words': 400}