Merge pull request #218 from wepay/include_kafka_key

Enable KCBQ to consume and stream Kafka Key into BigQuery.

Author: Bingqin Zhou

kcbq-api/src/main/java/com/wepay/kafka/connect/bigquery/api/KafkaSchemaRecordType.java

@@ -0,0 +1,21 @@
+package com.wepay.kafka.connect.bigquery.api;
+
+
+/**
+ *  Enum class for Kafka schema or record type, either value or key.
+ */
+public enum KafkaSchemaRecordType {
+
+    VALUE("value"),
+    KEY("key");
+
+    private final String str;
+
+    KafkaSchemaRecordType(String str) {
+        this.str = str;
+    }
+
+    public String toString() {
+        return this.str;
+    }
+}

kcbq-api/src/main/java/com/wepay/kafka/connect/bigquery/api/SchemaRetriever.java

@@ -22,9 +22,10 @@ public interface SchemaRetriever {
    * Retrieve the most current schema for the given topic.
    * @param table The table that will be created.
    * @param topic The topic to retrieve a schema for.
+   * @param schemaType The type of kafka schema, either "value" or "key".
    * @return The Schema for the given table.
    */
-  public Schema retrieveSchema(TableId table, String topic);
+  public Schema retrieveSchema(TableId table, String topic, KafkaSchemaRecordType schemaType);
 
   /**
    * Set the last seen schema for a given topic

kcbq-confluent/src/main/java/com/wepay/kafka/connect/bigquery/schemaregistry/schemaretriever/SchemaRegistrySchemaRetriever.java

@@ -2,6 +2,7 @@
 
 import com.google.cloud.bigquery.TableId;
 
+import com.wepay.kafka.connect.bigquery.api.KafkaSchemaRecordType;
 import com.wepay.kafka.connect.bigquery.api.SchemaRetriever;
 
 import io.confluent.connect.avro.AvroData;
@@ -61,8 +62,8 @@ public void configure(Map<String, String> properties) {
   }
 
   @Override
-  public Schema retrieveSchema(TableId table, String topic) {
-    String subject = getSubject(topic);
+  public Schema retrieveSchema(TableId table, String topic, KafkaSchemaRecordType schemaType) {
+    String subject = getSubject(topic, schemaType);
     try {
       logger.debug("Retrieving schema information for topic {} with subject {}", topic, subject);
       SchemaMetadata latestSchemaMetadata = schemaRegistryClient.getLatestSchemaMetadata(subject);
@@ -80,7 +81,7 @@ public Schema retrieveSchema(TableId table, String topic) {
   @Override
   public void setLastSeenSchema(TableId table, String topic, Schema schema) { }
 
-  private String getSubject(String topic) {
-    return topic + "-value";
+  private String getSubject(String topic, KafkaSchemaRecordType schemaType) {
+    return topic + "-" + schemaType.toString();
   }
 }

kcbq-confluent/src/test/java/com/wepay/kafka/connect/bigquery/schemaregistry/schemaretriever/SchemaRegistrySchemaRetrieverTest.java

@@ -7,6 +7,7 @@
 
 import com.google.cloud.bigquery.TableId;
 
+import com.wepay.kafka.connect.bigquery.api.KafkaSchemaRecordType;
 import io.confluent.connect.avro.AvroData;
 
 import io.confluent.kafka.schemaregistry.client.SchemaMetadata;
@@ -22,15 +23,17 @@ public class SchemaRegistrySchemaRetrieverTest {
   public void testRetrieveSchema() throws Exception {
     final TableId table = TableId.of("test", "kafka_topic");
     final String testTopic = "kafka-topic";
-    final String testSubject = "kafka-topic-value";
+    final String testSubjectValue = "kafka-topic-value";
+    final String testSubjectKey = "kafka-topic-key";
     final String testAvroSchemaString =
         "{\"type\": \"record\", "
         + "\"name\": \"testrecord\", "
         + "\"fields\": [{\"name\": \"f1\", \"type\": \"string\"}]}";
     final SchemaMetadata testSchemaMetadata = new SchemaMetadata(1, 1, testAvroSchemaString);
 
     SchemaRegistryClient schemaRegistryClient = mock(SchemaRegistryClient.class);
-    when(schemaRegistryClient.getLatestSchemaMetadata(testSubject)).thenReturn(testSchemaMetadata);
+    when(schemaRegistryClient.getLatestSchemaMetadata(testSubjectValue)).thenReturn(testSchemaMetadata);
+    when(schemaRegistryClient.getLatestSchemaMetadata(testSubjectKey)).thenReturn(testSchemaMetadata);
 
     SchemaRegistrySchemaRetriever testSchemaRetriever = new SchemaRegistrySchemaRetriever(
         schemaRegistryClient,
@@ -40,6 +43,7 @@ public void testRetrieveSchema() throws Exception {
     Schema expectedKafkaConnectSchema =
         SchemaBuilder.struct().field("f1", Schema.STRING_SCHEMA).name("testrecord").build();
 
-    assertEquals(expectedKafkaConnectSchema, testSchemaRetriever.retrieveSchema(table, testTopic));
+    assertEquals(expectedKafkaConnectSchema, testSchemaRetriever.retrieveSchema(table, testTopic, KafkaSchemaRecordType.VALUE));
+    assertEquals(expectedKafkaConnectSchema, testSchemaRetriever.retrieveSchema(table, testTopic, KafkaSchemaRecordType.KEY));
   }
 }

kcbq-connector/src/main/java/com/wepay/kafka/connect/bigquery/BigQuerySinkConnector.java

@@ -47,6 +47,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 
 /**
  * A {@link SinkConnector} used to delegate BigQuery data writes to
@@ -102,7 +103,9 @@ private SchemaManager getSchemaManager(BigQuery bigQuery) {
     }
     SchemaRetriever schemaRetriever = config.getSchemaRetriever();
     SchemaConverter<com.google.cloud.bigquery.Schema> schemaConverter = config.getSchemaConverter();
-    return new SchemaManager(schemaRetriever, schemaConverter, bigQuery);
+    Optional<String> kafkaKeyFieldName = config.getKafkaKeyFieldName();
+    Optional<String> kafkaDataFieldName = config.getKafkaDataFieldName();
+    return new SchemaManager(schemaRetriever, schemaConverter, bigQuery, kafkaKeyFieldName, kafkaDataFieldName);
   }
 
   private void ensureExistingTables(

kcbq-connector/src/main/java/com/wepay/kafka/connect/bigquery/BigQuerySinkTask.java

@@ -25,9 +25,11 @@
 import com.google.cloud.storage.BucketInfo;
 import com.google.cloud.storage.Storage;
 import com.google.common.annotations.VisibleForTesting;
+import com.wepay.kafka.connect.bigquery.api.KafkaSchemaRecordType;
 import com.wepay.kafka.connect.bigquery.api.SchemaRetriever;
 import com.wepay.kafka.connect.bigquery.config.BigQuerySinkConfig;
 import com.wepay.kafka.connect.bigquery.config.BigQuerySinkTaskConfig;
+import com.wepay.kafka.connect.bigquery.convert.KafkaDataBuilder;
 import com.wepay.kafka.connect.bigquery.convert.RecordConverter;
 import com.wepay.kafka.connect.bigquery.convert.SchemaConverter;
 import com.wepay.kafka.connect.bigquery.exception.SinkConfigConnectException;
@@ -63,6 +65,7 @@
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
+import java.util.Optional;
 
 /**
  * A {@link SinkTask} used to translate Kafka Connect {@link SinkRecord SinkRecords} into BigQuery
@@ -150,11 +153,18 @@ private PartitionedTableId getRecordTable(SinkRecord record) {
   }
 
   private RowToInsert getRecordRow(SinkRecord record) {
-    Map<String,Object> convertedRecord = recordConverter.convertRecord(record);
+    Map<String, Object> convertedRecord = recordConverter.convertRecord(record, KafkaSchemaRecordType.VALUE);
+    Optional<String> kafkaKeyFieldName = config.getKafkaKeyFieldName();
+    if (kafkaKeyFieldName.isPresent()) {
+      convertedRecord.put(kafkaKeyFieldName.get(), recordConverter.convertRecord(record, KafkaSchemaRecordType.KEY));
+    }
+    Optional<String> kafkaDataFieldName = config.getKafkaDataFieldName();
+    if (kafkaDataFieldName.isPresent()) {
+      convertedRecord.put(kafkaDataFieldName.get(), KafkaDataBuilder.buildKafkaDataRecord(record));
+    }
     if (config.getBoolean(config.SANITIZE_FIELD_NAME_CONFIG)) {
       convertedRecord = FieldNameSanitizer.replaceInvalidKeys(convertedRecord);
     }
-
     return RowToInsert.of(getRowId(record), convertedRecord);
   }
 
@@ -241,7 +251,9 @@ private SchemaManager getSchemaManager(BigQuery bigQuery) {
     schemaRetriever = config.getSchemaRetriever();
     SchemaConverter<com.google.cloud.bigquery.Schema> schemaConverter =
         config.getSchemaConverter();
-    return new SchemaManager(schemaRetriever, schemaConverter, bigQuery);
+    Optional<String> kafkaKeyFieldName = config.getKafkaKeyFieldName();
+    Optional<String> kafkaDataFieldName = config.getKafkaDataFieldName();
+    return new SchemaManager(schemaRetriever, schemaConverter, bigQuery, kafkaKeyFieldName, kafkaDataFieldName);
   }
 
   private BigQueryWriter getBigQueryWriter() {

kcbq-connector/src/main/java/com/wepay/kafka/connect/bigquery/SchemaManager.java

@@ -1,19 +1,26 @@
 package com.wepay.kafka.connect.bigquery;
 
+
 import com.google.cloud.bigquery.BigQuery;
+import com.google.cloud.bigquery.Field;
+import com.google.cloud.bigquery.LegacySQLTypeName;
 import com.google.cloud.bigquery.StandardTableDefinition;
 import com.google.cloud.bigquery.TableId;
 import com.google.cloud.bigquery.TableInfo;
 import com.google.cloud.bigquery.TimePartitioning;
-
+import com.wepay.kafka.connect.bigquery.api.KafkaSchemaRecordType;
 import com.wepay.kafka.connect.bigquery.api.SchemaRetriever;
+import com.wepay.kafka.connect.bigquery.convert.KafkaDataBuilder;
 import com.wepay.kafka.connect.bigquery.convert.SchemaConverter;
 
-import com.wepay.kafka.connect.bigquery.write.row.AdaptiveBigQueryWriter;
 import org.apache.kafka.connect.data.Schema;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Optional;
+
 /**
  * Class for managing Schemas of BigQuery tables (creating and updating).
  */
@@ -23,6 +30,8 @@ public class SchemaManager {
   private final SchemaRetriever schemaRetriever;
   private final SchemaConverter<com.google.cloud.bigquery.Schema> schemaConverter;
   private final BigQuery bigQuery;
+  private final Optional<String> kafkaKeyFieldName;
+  private final Optional<String> kafkaDataFieldName;
 
   /**
    * @param schemaRetriever Used to determine the Kafka Connect Schema that should be used for a
@@ -33,10 +42,14 @@ public class SchemaManager {
   public SchemaManager(
       SchemaRetriever schemaRetriever,
       SchemaConverter<com.google.cloud.bigquery.Schema> schemaConverter,
-      BigQuery bigQuery) {
+      BigQuery bigQuery,
+      Optional<String> kafkaKeyFieldName,
+      Optional<String> kafkaDataFieldName) {
     this.schemaRetriever = schemaRetriever;
     this.schemaConverter = schemaConverter;
     this.bigQuery = bigQuery;
+    this.kafkaKeyFieldName = kafkaKeyFieldName;
+    this.kafkaDataFieldName = kafkaDataFieldName;
   }
 
   /**
@@ -45,8 +58,9 @@ public SchemaManager(
    * @param topic The Kafka topic used to determine the schema.
    */
   public void createTable(TableId table, String topic) {
-    Schema kafkaConnectSchema = schemaRetriever.retrieveSchema(table, topic);
-    bigQuery.create(constructTableInfo(table, kafkaConnectSchema));
+    Schema kafkaValueSchema = schemaRetriever.retrieveSchema(table, topic, KafkaSchemaRecordType.VALUE);
+    Schema kafkaKeySchema = kafkaKeyFieldName.isPresent() ? schemaRetriever.retrieveSchema(table, topic, KafkaSchemaRecordType.KEY) : null;
+    bigQuery.create(constructTableInfo(table, kafkaKeySchema, kafkaValueSchema));
   }
 
   /**
@@ -55,26 +69,44 @@ public void createTable(TableId table, String topic) {
    * @param topic The Kafka topic used to determine the schema.
    */
   public void updateSchema(TableId table, String topic) {
-    Schema kafkaConnectSchema = schemaRetriever.retrieveSchema(table, topic);
-    TableInfo tableInfo = constructTableInfo(table, kafkaConnectSchema);
+    Schema kafkaValueSchema = schemaRetriever.retrieveSchema(table, topic, KafkaSchemaRecordType.VALUE);
+    Schema kafkaKeySchema = kafkaKeyFieldName.isPresent() ? schemaRetriever.retrieveSchema(table, topic, KafkaSchemaRecordType.KEY) : null;
+    TableInfo tableInfo = constructTableInfo(table, kafkaKeySchema, kafkaValueSchema);
     logger.info("Attempting to update table `{}` with schema {}",
         table, tableInfo.getDefinition().getSchema());
     bigQuery.update(tableInfo);
   }
 
   // package private for testing.
-  TableInfo constructTableInfo(TableId table, Schema kafkaConnectSchema) {
-    com.google.cloud.bigquery.Schema bigQuerySchema =
-        schemaConverter.convertSchema(kafkaConnectSchema);
+  TableInfo constructTableInfo(TableId table, Schema kafkaKeySchema, Schema kafkaValueSchema) {
+    com.google.cloud.bigquery.Schema bigQuerySchema = getBigQuerySchema(kafkaKeySchema, kafkaValueSchema);
     StandardTableDefinition tableDefinition = StandardTableDefinition.newBuilder()
         .setSchema(bigQuerySchema)
         .setTimePartitioning(TimePartitioning.of(TimePartitioning.Type.DAY))
         .build();
     TableInfo.Builder tableInfoBuilder =
         TableInfo.newBuilder(table, tableDefinition);
-    if (kafkaConnectSchema.doc() != null) {
-      tableInfoBuilder.setDescription(kafkaConnectSchema.doc());
+    if (kafkaValueSchema.doc() != null) {
+      tableInfoBuilder.setDescription(kafkaValueSchema.doc());
     }
     return tableInfoBuilder.build();
   }
+
+  private com.google.cloud.bigquery.Schema getBigQuerySchema(Schema kafkaKeySchema, Schema kafkaValueSchema) {
+      List<Field> allFields = new ArrayList<> ();
+      com.google.cloud.bigquery.Schema valueSchema = schemaConverter.convertSchema(kafkaValueSchema);
+      allFields.addAll(valueSchema.getFields());
+      if (kafkaKeyFieldName.isPresent()) {
+          com.google.cloud.bigquery.Schema keySchema = schemaConverter.convertSchema(kafkaKeySchema);
+          Field kafkaKeyField = Field.newBuilder(kafkaKeyFieldName.get(), LegacySQLTypeName.RECORD, keySchema.getFields())
+                  .setMode(Field.Mode.NULLABLE).build();
+          allFields.add(kafkaKeyField);
+      }
+      if (kafkaDataFieldName.isPresent()) {
+          Field kafkaDataField = KafkaDataBuilder.buildKafkaDataField(kafkaDataFieldName.get());
+          allFields.add(kafkaDataField);
+      }
+      return com.google.cloud.bigquery.Schema.of(allFields);
+  }
+
 }