Fixed #16 Added option to configure mapping of fields in SinkRecord to CQL columns

Author: aaruna

README.md

@@ -84,7 +84,8 @@ All the others (BLOB, INET, UUID, TIMEUUID, LIST, SET, MAP, CUSTOM, UDT, TUPLE,
 
 ## CassandraSink
 It stores Kafka SinkRecord in Cassandra tables. Currently, we only support STRUCT type in the SinkRecord. 
-The STRUCT can have multiple fields with primitive fieldtypes. We assume one-to-one mapping between the column names in the Cassandra sink table and the field names.
+The STRUCT can have multiple fields with primitive fieldtypes. 
+By default, we assume one-to-one mapping between the column names in the Cassandra sink table and the field names.
 
 Say, the SinkRecords has the following STRUCT value
 ```
@@ -97,6 +98,23 @@ Say, the SinkRecords has the following STRUCT value
 
 Then the Cassandra table should have the columns - id, username, text  
 
+We also support specifying the field name mapping to column names, using the property `cassandra.sink.field.mapping`
+Say, the SinkRecords has the following STRUCT value
+```
+{
+    'id': 1,
+    'user': {
+        'id': 123,
+        'name': 'Foo',
+        'email': 'foo@bar.com'
+    },
+    'text': 'This is my first tweet'
+}
+``` 
+and the `cassandra.sink.field.mapping` has the value `{'id': 'id', 'user': {'id': 'uid', 'name': 'username'}, 'text': 'tweet_text'}`
+Then the Cassandra table should have the columns - id, uid, username, tweet_text. 
+Note that since we did not specify any mapping for 'user.email', it is ignored and not inserted in the Cassandra Sink table.
+
 Note: The library does not create the Cassandra tables - users are expected to create those before starting the sink
 
 ## Configuration
@@ -132,6 +150,7 @@ Refer `examples/config` for sample configuration files
 |--------           |----------------------------|-----------------------|
 | cassandra.sink.route.\<topic_name\> | The table to write the SinkRecords to, \<keyspace\>.\<tableName\> | |
 | cassandra.sink.consistency | The consistency level for writes to Cassandra. | LOCAL_QUORUM |
+| cassandra.sink.field.mapping | The JSON String mapping field names to column names. | |
 
 
 ## Building from Source

build.sbt

@@ -33,14 +33,16 @@ libraryDependencies ++= Seq(
   "org.joda"               %  "joda-convert"          % "1.8.1",
   "org.scalatest"          %% "scalatest"             % "2.2.6"       % "test,it",
   "org.mockito"            % "mockito-core"           % "2.0.34-beta" % "test,it",
-  "ch.qos.logback"         % "logback-classic"        % "1.1.3"       % "test,it",
-  CrossVersion.partialVersion(scalaVersion.value) match {
-    case Some((2, minor)) if minor < 11 =>
-      "org.slf4j"                  % "slf4j-api"      % "1.7.13"
-    case _ =>
-      "com.typesafe.scala-logging" %% "scala-logging" % "3.1.0"
-  }
-)
+  "ch.qos.logback"         % "logback-classic"        % "1.1.3"       % "test,it"
+) ++ (CrossVersion.partialVersion(scalaVersion.value) match {
+  case Some((2, minor)) if minor < 11 => Seq(
+    "org.slf4j"                  % "slf4j-api"     % "1.7.13"
+  )
+  case _ => Seq(
+    "com.typesafe.scala-logging" %% "scala-logging" % "3.1.0",
+    "org.scala-lang.modules" %% "scala-parser-combinators" % "1.0.4"
+  )
+})
 
 publishMavenStyle := true
 

src/it/resources/setup.cql

@@ -26,6 +26,13 @@ CREATE TABLE IF NOT EXISTS test.kv (
   value int,
   PRIMARY KEY (key));
 
+CREATE TABLE IF NOT EXISTS test.fieldmap (
+  new_key text,
+  new_value int,
+  new_nested text,
+  new_dnested text,
+  PRIMARY KEY (new_key));
+
 CREATE TABLE test.playlists (
   id bigint,
   song_order int,

src/it/scala/com/tuplejump/kafka/connect/cassandra/CassandraSinkTaskSpec.scala

@@ -19,16 +19,17 @@
 package com.tuplejump.kafka.connect.cassandra
 
 import scala.collection.JavaConverters._
+import scala.util.parsing.json.JSONObject
 import org.apache.kafka.connect.data.{Schema, SchemaBuilder, Struct}
 import org.apache.kafka.connect.sink.{SinkRecord, SinkTaskContext}
 
 class CassandraSinkTaskSpec extends AbstractFlatSpec {
 
-  val topicName = "test_kv_topic"
-  val tableName = "test.kv"
-  val config = sinkProperties(Map(topicName -> tableName))
-
   it should "start sink task" in {
+    val topicName = "test_kv_topic"
+    val tableName = "test.kv"
+    val config = sinkProperties(Map(topicName -> tableName))
+
     val sinkTask = new CassandraSinkTask()
     val mockContext = mock[SinkTaskContext]
 
@@ -38,6 +39,10 @@ class CassandraSinkTaskSpec extends AbstractFlatSpec {
   }
 
   it should "save records in cassandra" in {
+    val topicName = "test_kv_topic"
+    val tableName = "test.kv"
+    val config = sinkProperties(Map(topicName -> tableName))
+
     val sinkTask = new CassandraSinkTask()
     val mockContext = mock[SinkTaskContext]
 
@@ -64,5 +69,76 @@ class CassandraSinkTaskSpec extends AbstractFlatSpec {
     rowCount should be(2)
     cc.shutdown()
   }
+
+
+  it should "save records in cassandra with custom field mapping" in {
+    val topicName = "test_fieldmap_topic"
+    val tableName = "test.fieldmap"
+    val config = sinkProperties(Map(topicName -> tableName))
+
+    val sinkTask = new CassandraSinkTask()
+    val mockContext = mock[SinkTaskContext]
+
+    val fieldMapping: JSONObject = JSONObject(Map(
+      "key" -> "new_key",
+      "value" -> "new_value",
+      "nvalue" -> JSONObject(Map(
+        "blah1" -> "new_nested",
+        "blah2" -> JSONObject(Map(
+          "blah2" -> "new_dnested"
+        ))
+      ))
+    ))
+
+    sinkTask.initialize(mockContext)
+    sinkTask.start((config + ("cassandra.sink.field.mapping" -> fieldMapping.toString())).asJava)
+
+    val doubleNestedSchema = SchemaBuilder.struct.name("dnestedSchema").version(1)
+      .field("blah1", Schema.STRING_SCHEMA)
+      .field("blah2", Schema.STRING_SCHEMA).build
+    val nestedSchema = SchemaBuilder.struct.name("nestedSchema").version(1)
+      .field("blah1", Schema.STRING_SCHEMA)
+      .field("blah2", doubleNestedSchema).build
+    val valueSchema = SchemaBuilder.struct.name("record").version(1)
+      .field("key", Schema.STRING_SCHEMA)
+      .field("value", Schema.INT32_SCHEMA)
+      .field("nvalue", nestedSchema).build
+
+    val dnestedValue1 = new Struct(doubleNestedSchema)
+      .put("blah1", "dnes_blah1_1")
+      .put("blah2", "dnes_blah2_1")
+    val nestedValue1 = new Struct(nestedSchema)
+      .put("blah1", "nes_blah1_1")
+      .put("blah2", dnestedValue1)
+    val value1 = new Struct(valueSchema)
+      .put("key", "pqr")
+      .put("value", 15)
+      .put("nvalue", nestedValue1)
+
+    val dnestedValue2 = new Struct(doubleNestedSchema)
+      .put("blah1", "dnes_blah1_2")
+      .put("blah2", "dnes_blah2_2")
+    val nestedValue2 = new Struct(nestedSchema)
+      .put("blah1", "nes_blah1_2")
+      .put("blah2", dnestedValue2)
+    val value2 = new Struct(valueSchema)
+      .put("key", "abc")
+      .put("value", 17)
+      .put("nvalue", nestedValue2)
+
+    val record1 = new SinkRecord(topicName, 1, SchemaBuilder.struct.build, "key", valueSchema, value1, 0)
+    val record2 = new SinkRecord(topicName, 1, SchemaBuilder.struct.build, "key", valueSchema, value2, 0)
+
+    sinkTask.put(List(record1, record2).asJavaCollection)
+
+    sinkTask.stop()
+
+    val cc = CassandraCluster.local
+    val session = cc.session
+    val result = session.execute(s"select count(1) from $tableName").one()
+    val rowCount = result.getLong(0)
+    rowCount should be(2)
+    cc.shutdown()
+  }
 }
 

src/main/scala/com/tuplejump/kafka/connect/cassandra/CassandraSinkTask.scala

@@ -52,7 +52,7 @@ class CassandraSinkTask extends SinkTask with CassandraTask {
   private def write(sc: SinkConfig, byTopic: Iterable[SinkRecord]): Unit = {
     // TODO needs ticket: if (byTopic.size > 1) boundWrite(sc, byTopic) else
       for (record <- byTopic) {
-        val query = record.as(sc.schema.namespace)
+        val query = record.as(sc.schema.namespace, sc.options.fieldMapping)
         Try(session.executeAsync(query.cql)) recover { case NonFatal(e) =>
           throw new ConnectException(
             s"Error executing ${byTopic.size} records for schema '${sc.schema}'.", e)
@@ -64,7 +64,7 @@ class CassandraSinkTask extends SinkTask with CassandraTask {
   private def boundWrite(sc: SinkConfig, byTopic: Iterable[SinkRecord]): Unit = {
     val statement = prepare(session, sc)
     val futures = for (record <- byTopic) yield {
-      val query = record.as(sc.schema.namespace)
+      val query = record.as(sc.schema.namespace, sc.options.fieldMapping)
       try {
         val bs = statement.bind(query.cql)
         session.executeAsync(bs)

src/main/scala/com/tuplejump/kafka/connect/cassandra/ConnectorSyntax.scala

@@ -107,9 +107,12 @@ private[cassandra] object Syntax {
       namespace.length >= 3 || namespace.contains(".")
     }
 
-    def apply(namespace: String, columnNames: List[ColumnName], columnValues: String): SinkQuery = {
-      val columns = columnNames.mkString(",")
-      SinkQuery(s"INSERT INTO $namespace($columns) VALUES($columnValues)")
+    def apply(namespace: String, columnNamesVsValues: Map[ColumnName, String]): SinkQuery = {
+      val query = columnNamesVsValues.view.map(e => Vector(e._1, e._2)).transpose match {
+        case columnNames :: columnValues :: Nil =>
+          s"INSERT INTO ${namespace}(${columnNames.mkString(",")}) VALUES(${columnValues.mkString(",")})"
+      }
+      SinkQuery(query)
     }
   }
 

src/main/scala/com/tuplejump/kafka/connect/cassandra/TaskConfig.scala

@@ -18,8 +18,10 @@ package com.tuplejump.kafka.connect.cassandra
 
 import scala.collection.immutable
 import scala.util.control.NonFatal
+import scala.util.parsing.json.JSON
 import org.apache.kafka.common.config.ConfigException
 import org.apache.kafka.connect.connector.Task
+import org.apache.kafka.connect.errors.DataException
 import com.datastax.driver.core.{TableMetadata, ConsistencyLevel}
 import InternalConfig._
 
@@ -123,6 +125,9 @@ object TaskConfig {
   final val SinkConsistency: Key = "cassandra.sink.consistency"
   final val DefaultSinkConsistency = ConsistencyLevel.LOCAL_QUORUM
 
+  final val FieldMapping: Key = "cassandra.sink.field.mapping"
+  final val DefaultFieldMapping = Map.empty[String, String]
+
   /* **** Task config **** */
   final val TaskParallelismStrategy: Key = "cassandra.task.parallelism"
 
@@ -156,6 +161,10 @@ private[cassandra] object InternalConfig {
   def toInt(a: String): Int = a.toInt
   def toLong(a: String): Long = a.toLong
   def toConsistency(a: String): ConsistencyLevel = ConsistencyLevel.valueOf(a)
+  def toMap(a: String): Map[String, Any] = JSON.parseFull(a) collect {
+    case data: Map[_, _] => data.asInstanceOf[Map[String, Any]]
+  } getOrElse(throw new DataException(s"Field mapping type for '$a' is not supported."))
+
 
   /** A Cassandra `keyspace.table` to Kafka topic mapping.
     *
@@ -319,15 +328,21 @@ private[cassandra] object InternalConfig {
   sealed trait ClusterQueryOptions
 
   /** Settings related for individual queries, can be set per keyspace.table. */
-  final case class WriteOptions(consistency: ConsistencyLevel) extends ClusterQueryOptions
+  final case class WriteOptions(consistency: ConsistencyLevel,
+                                fieldMapping: Map[String, Any]) extends ClusterQueryOptions
 
   object WriteOptions {
 
-    def apply(config: Map[String,String]): WriteOptions =
-      WriteOptions(config.valueOr[ConsistencyLevel](
-        SinkConsistency, toConsistency, DefaultSourceConsistency))
+    def apply(config: Map[String, String]): WriteOptions = {
+      WriteOptions(
+        consistency = config.valueOr[ConsistencyLevel](
+          SinkConsistency, toConsistency, DefaultSourceConsistency),
+        fieldMapping = config.valueOr[Map[String, Any]](
+          FieldMapping, toMap, DefaultFieldMapping
+        )
+      )
+    }
   }
-
   /** Settings related for individual queries, can be set per keyspace.table. */
   final case class ReadOptions(splitSize: Int,
                                fetchSize: Int,

src/main/scala/com/tuplejump/kafka/connect/cassandra/package.scala

@@ -16,6 +16,8 @@
 
 package com.tuplejump.kafka.connect
 
+import org.apache.kafka.connect.data.Field
+
 /** Common package operations. */
 package object cassandra {
   import java.util.{List => JList, Map => JMap, Date => JDate}
@@ -69,29 +71,90 @@ package object cassandra {
 
   implicit class SinkRecordOps(record: SinkRecord) {
 
-    def as(namespace: String): SinkQuery = {
+    def as(namespace: String, fieldMapping: Map[String, Any]): SinkQuery = {
+      val colNamesVsValues: Map[String, String] = {
+        if (fieldMapping.isEmpty) {
+          toCqlData
+        } else {
+          toCqlData(fieldMapping)
+        }
+      }
+      SinkQuery(namespace, colNamesVsValues)
+    }
+
+    def toCqlData(): (Map[String, String]) = {
       val schema = record.valueSchema
-      val columnNames = schema.asColumnNames
-      val columnValues = schema.`type`() match {
+      schema.`type`() match {
         case STRUCT =>
-          val struct: Struct = record.value.asInstanceOf[Struct]
-          columnNames.map(convert(schema, struct, _)).mkString(",")
-        case other => throw new DataException(
-          s"Unable to create insert statement with unsupported value schema type $other.")
+          schema.fields.asScala.map { field =>
+            field.name -> convert(schema, record.value.asInstanceOf[Struct], field)
+          }.toMap
+        case other =>
+          throw new DataException(
+            s"Unable to create insert statement with unsupported value schema type $other.")
+      }
+    }
+
+    def toCqlData(fieldMapping: Map[String, Any]): Map[String, String] = {
+      record.valueSchema.`type`() match {
+        case STRUCT =>
+          toColNamesVsValues(Map.empty[String, String], record.value.asInstanceOf[Struct], fieldMapping)
+        case other =>
+          throw new DataException(
+            s"Unable to create insert statement with unsupported value schema type $other.")
+      }
+    }
+
+    // scalastyle:off
+    private def toColNamesVsValues(colNameVsValues: Map[String, String],
+                                   struct: Struct, fieldMapping: Map[String, Any]): Map[String, String] = {
+      lazy val exception = new DataException(s"Mismatch between fieldMapping and Schema")
+      var result: Map[String, String] = colNameVsValues
+      struct.schema.fields.asScala.foreach { field =>
+        val fieldMappingValue = fieldMapping.get(field.name)
+        field.schema.`type`() match {
+          case STRUCT =>
+            fieldMappingValue match {
+              case Some(value) =>
+                value match {
+                  case newMap: Map[_, _] =>
+                    result = toColNamesVsValues(
+                      result,
+                      struct.get(field).asInstanceOf[Struct],
+                      newMap.asInstanceOf[Map[String, Any]]
+                    )
+                  case _ =>
+                    throw exception
+                }
+              case None =>
+            }
+          case _ =>
+            fieldMappingValue match {
+              case Some(value) =>
+                value match {
+                  case strValue: String =>
+                    result += (strValue -> convert(field.schema, struct, field))
+                  case _ =>
+                    throw exception
+                }
+              case None =>
+            }
+        }
       }
-      SinkQuery(namespace, columnNames, columnValues)
+      result
     }
+    // scalastyle:on
 
     /* TODO support all types. */
-    def convert(schema: Schema, result: Struct, col: String): AnyRef =
-      schema.field(col).schema match {
+    def convert(schema: Schema, result: Struct, field: Field): String =
+      field.schema match {
         case x if x.`type`() == Schema.STRING_SCHEMA.`type`() =>
-          s"'${result.get(col).toString}'"
+          s"'${result.get(field).toString}'"
         case x if x.name() == Timestamp.LOGICAL_NAME =>
-          val time = Timestamp.fromLogical(x, result.get(col).asInstanceOf[JDate])
+          val time = Timestamp.fromLogical(x, result.get(field).asInstanceOf[JDate])
           s"$time"
         case y =>
-          result.get(col)
+          String.valueOf(result.get(field))
       }
 
     def asColumnNames: List[ColumnName] =