com.ubirch.kafka.package.scala Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of ubirch-kafka-envelope Show documentation
kafka message envelope and utils
The newest version!
package com.ubirch

import java.io.ByteArrayInputStream
import java.lang.{Iterable => JIterable}
import java.nio.charset.StandardCharsets.UTF_8
import java.util
import java.util.Optional

import com.fasterxml.jackson.databind.node._
import com.fasterxml.jackson.databind.{DeserializationFeature, JsonNode, MapperFeature, ObjectMapper, SerializationFeature}
import com.ubirch.protocol.ProtocolMessage
import org.apache.kafka.clients.consumer.ConsumerRecord
import org.apache.kafka.clients.producer.ProducerRecord
import org.apache.kafka.common.header.internals.{RecordHeader, RecordHeaders}
import org.apache.kafka.common.header.{Header, Headers}
import org.apache.kafka.common.record.TimestampType
import org.apache.kafka.common.serialization.{Deserializer, Serializer}
import org.json4s.jackson.JsonMethods._
import org.json4s.jackson.Serialization._
import org.json4s.{CustomSerializer, DefaultFormats, Formats}

import scala.collection.JavaConverters._
import scala.collection.breakOut

package object kafka {
  implicit val formats: Formats = DefaultFormats + Json4sProtocolMessageSerializer

  object EnvelopeSerializer extends Serializer[MessageEnvelope] {
    override def configure(configs: util.Map[String, _], isKey: Boolean): Unit = {}

    override def close(): Unit = {}

    def serializeToString(data: MessageEnvelope): String = {
      write(data)
    }

    override def serialize(_topic: String, data: MessageEnvelope): Array[Byte] = {
      write(data).getBytes(UTF_8)
    }
  }

  object EnvelopeDeserializer extends Deserializer[MessageEnvelope] {
    override def configure(configs: util.Map[String, _], isKey: Boolean): Unit = {}

    override def close(): Unit = {}

    override def deserialize(_topic: String, data: Array[Byte]): MessageEnvelope = {
      read[MessageEnvelope](new ByteArrayInputStream(data))
    }
  }

  implicit class RichAnyConsumerRecord[K, V](val cr: ConsumerRecord[K, V]) extends AnyVal {
    def copy[K1, V1](topic: String = cr.topic(),
                     partition: Int = cr.partition(),
                     offset: Long = cr.offset(),
                     timestamp: Long = cr.timestamp(),
                     timestampType: TimestampType = cr.timestampType(),
                     checksum: Long = cr.checksum(),
                     serializedKeySize: Int = cr.serializedKeySize(),
                     serializedValueSize: Int = cr.serializedValueSize(),
                     key: K1 = cr.key(),
                     value: V1 = cr.value(),
                     headers: Headers = cr.headers(),
                     leaderEpoch: Optional[Integer] = cr.leaderEpoch()) =
      new ConsumerRecord[K1, V1](topic, partition, offset, timestamp, timestampType, checksum, serializedKeySize,
        serializedValueSize, key, value, headers, leaderEpoch)

    def headersScala: Map[String, String] = cr.headers().asScala
      .map(h => h.key() -> new String(h.value(), UTF_8))(breakOut)

    def toProducerRecord[V1](topic: String, value: V1 = cr.value(), partition: Integer = null, headers: JIterable[Header] = cr.headers()): ProducerRecord[K, V1] = {
      new ProducerRecord(topic, partition, cr.key(), value, headers)
    }

    def withHeaders(headers: (String, String)*): ConsumerRecord[K, V] = {
      val headersIterable: JIterable[Header] = headers
        .map(p => new RecordHeader(p._1, p._2.getBytes(UTF_8)): Header).asJava
      copy(headers = new RecordHeaders(headersIterable))
    }

    def withExtraHeaders(extraHeaders: (String, String)*): ConsumerRecord[K, V] = {
      val allHeaders = headersScala.toList ++ extraHeaders
      withHeaders(allHeaders: _*)
    }
  }

  implicit class RichMessageEnvelopeConsumerRecord[K](val cr: ConsumerRecord[K, MessageEnvelope]) extends AnyVal {
    def withExtraContext[T](key: String, value: T)(implicit formats: Formats): ConsumerRecord[K, MessageEnvelope] = {
      val newValue = cr.value().withExtraContext(key, value)
      cr.copy(value = newValue)
    }

    def withExtraContext[T](fields: (String, Any)*)(implicit formats: Formats): ConsumerRecord[K, MessageEnvelope] = {
      val newValue = cr.value().withExtraContext(fields: _*)
      cr.copy(value = newValue)
    }
  }

  implicit class RichAnyProducerRecord[K, V](val pr: ProducerRecord[K, V]) extends AnyVal {
    def copy[K1, V1](topic: String = pr.topic(),
                     partition: Integer = pr.partition(),
                     timestamp: Long = pr.timestamp(),
                     key: K1 = pr.key(),
                     value: V1 = pr.value(),
                     headers: JIterable[Header] = pr.headers()) =
      new ProducerRecord[K1, V1](topic, partition, timestamp, key, value, headers)

    def headersScala: Map[String, String] = pr.headers().asScala
      .map(h => h.key() -> new String(h.value(), UTF_8))(breakOut)

    def withHeaders(headers: (String, String)*): ProducerRecord[K, V] = {
      val headersIterable: JIterable[Header] = headers
        .map(p => new RecordHeader(p._1, p._2.getBytes(UTF_8)): Header).asJava
      copy(headers = new RecordHeaders(headersIterable))
    }

    def withExtraHeaders(extraHeaders: (String, String)*): ProducerRecord[K, V] = {
      val allHeaders = headersScala.toList ++ extraHeaders
      withHeaders(allHeaders: _*)
    }
  }

  // based on JSONProtocolEncoder, but without messing with signatures
  private val protocolMessageMapper = {
    val mapper = new ObjectMapper()
    mapper.configure(MapperFeature.SORT_PROPERTIES_ALPHABETICALLY, true)
    mapper.configure(SerializationFeature.ORDER_MAP_ENTRIES_BY_KEYS, true)
    mapper.configure(DeserializationFeature.USE_BIG_INTEGER_FOR_INTS, true)
    mapper
  }

  private object Json4sProtocolMessageSerializer extends CustomSerializer[ProtocolMessage](_ => ( {
    case jVal =>
      protocolMessageMapper.treeToValue[ProtocolMessage](asJsonNode(jVal), classOf[ProtocolMessage])
  }, {
    case p: ProtocolMessage =>
      val tree: ObjectNode = protocolMessageMapper.valueToTree(p)

      // workaround for Json4s not handling jackson binary nodes
      // - this can be handled better in future versions of Json4s
      def replaceBinaryNodesWithTextNodes(tree: ContainerNode[_]): Unit = {
        def inner[K](withKeys: (K => Unit) => Unit, get: K => JsonNode, set: (K, JsonNode) => Unit): Unit = {
          withKeys { key =>
            get(key) match {
              case binary: BinaryNode => set(key, new TextNode(binary.asText()))
              case containerNode: ContainerNode[_] => replaceBinaryNodesWithTextNodes(containerNode)
              case _ => // do nothing
            }
          }
        }

        tree match {
          case o: ObjectNode => inner[String](o.fieldNames().asScala.foreach, o.get, o.set)
          case a: ArrayNode => inner[Int]((0 until a.size()).foreach, a.get, a.set)
        }
      }

      replaceBinaryNodesWithTextNodes(tree)
      fromJsonNode(tree)
  }))

}