com.github.fbascheper.kafka.connect.telegram.mapper.StructFieldsValuesExtractor.scala Maven / Gradle / Ivy

Go to download
package com.github.fbascheper.kafka.connect.telegram.mapper

/**
  * Created by [email protected] on 29/05/16.
  *
  * @since 20-09-2018
  */

import java.text.SimpleDateFormat
import java.util.TimeZone

import org.apache.kafka.connect.data._

import scala.collection.JavaConversions._

trait StructFieldsValuesExtractor {
  def get(struct: Struct): Seq[(String, Any)]
}

/**
  * Extracts fields from a SinkRecord Struct based on a specified set of provided columns.
  *
  * @param includeAllFields Boolean indicating if all the fields from the SinkRecord are to be written to the sink
  * @param fieldsAliasMap   A map of fields and their alias,if provided, to extract from the SinkRecord
  **/
case class StructFieldsExtractor(includeAllFields: Boolean, fieldsAliasMap: Map[String, String]) extends StructFieldsValuesExtractor {

  /**
    * Get a sequence of columns names to column values for a given struct
    *
    * @param struct A SinkRecord struct
    * @return a Sequence of column names and values
    **/
  def get(struct: Struct): Seq[(String, AnyRef)] = {
    val schema = struct.schema()
    val fields: Seq[Field] = if (includeAllFields) schema.fields()
    else schema.fields().filter(f => fieldsAliasMap.contains(f.name()))

    val fieldsAndValues = fields.flatMap(field =>
      getFieldValue(field, struct).map(value => fieldsAliasMap.getOrElse(field.name(), field.name()) -> value))
    fieldsAndValues
  }

  /**
    * For a field in a struct return the value
    *
    * @param field  A field to return the value for
    * @param struct A struct to extract the field from
    * @return an optional value for the field
    **/
  private def getFieldValue(field: Field, struct: Struct): Option[AnyRef] = {
    Option(struct.get(field)) match {
      case None => None
      case Some(value) =>
        val fieldName = field.name()
        Option(field.schema().name()).collect {
          case Decimal.LOGICAL_NAME =>
            value match {
              case bd: BigDecimal => bd
              case array: Array[Byte] => Decimal.toLogical(field.schema, value.asInstanceOf[Array[Byte]])
            }
          case Date.LOGICAL_NAME =>
            value.asInstanceOf[Any] match {
              case d: java.util.Date => d
              case i: Int => Date.toLogical(field.schema, i)
              case _ => throw new IllegalArgumentException(s"Can't convert $value to Date for schema:${field.schema().`type`()}")
            }
          case Time.LOGICAL_NAME =>
            value.asInstanceOf[Any] match {
              case i: Int => Time.toLogical(field.schema, value.asInstanceOf[Int])
              case d: java.util.Date => d
              case _ => throw new IllegalArgumentException(s"Can't convert $value to Date for schema:${field.schema().`type`()}")
            }
          case Timestamp.LOGICAL_NAME =>
            value.asInstanceOf[Any] match {
              case l: Long => Timestamp.toLogical(field.schema, l)
              case d: java.util.Date => d
              case _ => throw new IllegalArgumentException(s"Can't convert $value to Date for schema:${field.schema().`type`()}")
            }
        }.orElse {
          val v = field.schema().`type`() match {
            case Schema.Type.BOOLEAN => struct.getBoolean(fieldName)
            case Schema.Type.BYTES => struct.getBytes(fieldName)
            case Schema.Type.FLOAT32 => struct.getFloat32(fieldName)
            case Schema.Type.FLOAT64 => struct.getFloat64(fieldName)
            case Schema.Type.INT8 => struct.getInt8(fieldName)
            case Schema.Type.INT16 => struct.getInt16(fieldName)
            case Schema.Type.INT32 => struct.getInt32(fieldName)
            case Schema.Type.INT64 => struct.getInt64(fieldName)
            case Schema.Type.STRING => struct.getString(fieldName)
            case Schema.Type.ARRAY => struct.getArray(fieldName)
            case Schema.Type.MAP => struct.getMap(fieldName)
            case Schema.Type.STRUCT => struct.getStruct(fieldName)
            case other => sys.error(s"$other is not a recognized schema")
          }
          Some(v)
        }
    }
  }
}


object StructFieldsExtractor {
  val DateFormat = new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss.SSS'Z'")
  val TimeFormat: SimpleDateFormat = new SimpleDateFormat("HH:mm:ss.SSSZ")
  DateFormat.setTimeZone(TimeZone.getTimeZone("UTC"))
}