Refactor.

Author: Bingqin Zhou

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

@@ -22,10 +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 subjectType The type of kafka schema subject, either "value" or "key".
+   * @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, String subjectType);
+  public Schema retrieveSchema(TableId table, String topic, String schemaType);
 
   /**
    * Set the last seen schema for a given topic

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

@@ -61,8 +61,8 @@ public void configure(Map<String, String> properties) {
   }
 
   @Override
-  public Schema retrieveSchema(TableId table, String topic, String subjectType) {
-    String subject = getSubject(topic, subjectType);
+  public Schema retrieveSchema(TableId table, String topic, String 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 +80,7 @@ public Schema retrieveSchema(TableId table, String topic, String subjectType) {
   @Override
   public void setLastSeenSchema(TableId table, String topic, Schema schema) { }
 
-  private String getSubject(String topic, String subjectType) {
-    return topic + "-" + subjectType;
+  private String getSubject(String topic, String schemaType) {
+    return topic + "-" + schemaType;
   }
 }

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

@@ -29,6 +29,7 @@
 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.KafkaSchemaRecordType;
 import com.wepay.kafka.connect.bigquery.convert.RecordConverter;
 import com.wepay.kafka.connect.bigquery.convert.SchemaConverter;
 import com.wepay.kafka.connect.bigquery.exception.SinkConfigConnectException;
@@ -48,7 +49,6 @@
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.config.ConfigException;
 import org.apache.kafka.common.record.TimestampType;
-import org.apache.kafka.connect.data.Schema;
 import org.apache.kafka.connect.errors.ConnectException;
 import org.apache.kafka.connect.sink.SinkRecord;
 import org.apache.kafka.connect.sink.SinkTask;
@@ -152,9 +152,9 @@ private PartitionedTableId getRecordTable(SinkRecord record) {
   }
 
   private RowToInsert getRecordRow(SinkRecord record) {
-    Map<String, Object> convertedRecord = recordConverter.convertRecord(record.valueSchema(), record.value());
+    Map<String, Object> convertedRecord = recordConverter.convertRecord(record, KafkaSchemaRecordType.VALUE);
     if (config.getBoolean(config.INCLUDE_KAFKA_KEY_CONFIG)) {
-      convertedRecord.put(config.getString(config.KAFKA_KEY_FIELD_NAME_CONFIG), recordConverter.convertRecord(record.keySchema(), record.key()));
+      convertedRecord.put(config.getString(config.KAFKA_KEY_FIELD_NAME_CONFIG), recordConverter.convertRecord(record, KafkaSchemaRecordType.KEY));
     }
     if (config.getBoolean(config.INCLUDE_KAFKA_DATA_CONFIG)) {
       convertedRecord.put(config.getString(config.KAFKA_DATA_FIELD_NAME_CONFIG), KafkaDataBuilder.buildKafkaDataRecord(record));

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

@@ -4,6 +4,7 @@
 import com.google.cloud.bigquery.*;
 import com.wepay.kafka.connect.bigquery.api.SchemaRetriever;
 import com.wepay.kafka.connect.bigquery.convert.KafkaDataBuilder;
+import com.wepay.kafka.connect.bigquery.convert.KafkaSchemaRecordType;
 import com.wepay.kafka.connect.bigquery.convert.SchemaConverter;
 
 import org.apache.kafka.connect.data.Schema;
@@ -18,8 +19,6 @@
  */
 public class SchemaManager {
   private static final Logger logger = LoggerFactory.getLogger(SchemaManager.class);
-  private static final String VALUE = "value";
-  private static final String KEY = "key";
 
   private final SchemaRetriever schemaRetriever;
   private final SchemaConverter<com.google.cloud.bigquery.Schema> schemaConverter;
@@ -58,8 +57,8 @@ public SchemaManager(
    * @param topic The Kafka topic used to determine the schema.
    */
   public void createTable(TableId table, String topic) {
-    Schema kafkaValueSchema = schemaRetriever.retrieveSchema(table, topic, VALUE);
-    Schema kafkaKeySchema = includeKafkaKey ? schemaRetriever.retrieveSchema(table, topic, KEY) : null;
+    Schema kafkaValueSchema = schemaRetriever.retrieveSchema(table, topic, KafkaSchemaRecordType.VALUE.toString());
+    Schema kafkaKeySchema = includeKafkaKey ? schemaRetriever.retrieveSchema(table, topic, KafkaSchemaRecordType.KEY.toString()) : null;
     bigQuery.create(constructTableInfo(table, kafkaKeySchema, kafkaValueSchema));
   }
 
@@ -69,8 +68,8 @@ 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 kafkaValueSchema = schemaRetriever.retrieveSchema(table, topic, VALUE);
-    Schema kafkaKeySchema = includeKafkaKey ? schemaRetriever.retrieveSchema(table, topic, KEY) : null;
+    Schema kafkaValueSchema = schemaRetriever.retrieveSchema(table, topic, KafkaSchemaRecordType.VALUE.toString());
+    Schema kafkaKeySchema = includeKafkaKey ? schemaRetriever.retrieveSchema(table, topic, KafkaSchemaRecordType.KEY.toString()) : null;
     TableInfo tableInfo = constructTableInfo(table, kafkaKeySchema, kafkaValueSchema);
     logger.info("Attempting to update table `{}` with schema {}",
         table, tableInfo.getDefinition().getSchema());

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

@@ -65,15 +65,15 @@ public BigQueryRecordConverter(boolean shouldConvertDoubleSpecial) {
 
   /**
    * Convert a {@link SinkRecord} into the contents of a BigQuery {@link RowToInsert}.
-   * @param kafkaConnectSchema The schema of the kafka connect record to convert. Must be of type {@link Struct},
-   *                           in order to translate into a row format that requires each field to
-   *                           consist of both a name and a value.
-   * @param kafkaConnectStruct The struct of the kafka connect record to convert. Must be of type {@link Struct},
-   *                           in order to translate into a row format that requires each field to
-   *                           consist of both a name and a value.
+   * @param record The Kafka Connect record to convert. Must be of type {@link Struct},
+   *               in order to translate into a row format that requires each field to
+   *               consist of both a name and a value.
+   * @param recordType The type of the record to convert, either value or key.
    * @return The result BigQuery row content.
    */
-  public Map<String, Object> convertRecord(Schema kafkaConnectSchema, Object kafkaConnectStruct) {
+  public Map<String, Object> convertRecord(SinkRecord record, KafkaSchemaRecordType recordType) {
+    Schema kafkaConnectSchema = recordType == KafkaSchemaRecordType.KEY ? record.keySchema() : record.valueSchema();
+    Object kafkaConnectStruct = recordType == KafkaSchemaRecordType.KEY ? record.key() : record.value();
     if (kafkaConnectSchema == null) {
       if (kafkaConnectStruct instanceof Map) {
         return (Map<String, Object>) convertSchemalessRecord(kafkaConnectStruct);

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

@@ -0,0 +1,14 @@
+package com.wepay.kafka.connect.bigquery.convert;
+
+
+public enum KafkaSchemaRecordType {
+    VALUE("value"),
+    KEY("key");
+    String str;
+    KafkaSchemaRecordType(String str) {
+        this.str = str;
+    }
+    public String toString() {
+        return this.str;
+    }
+}

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

@@ -18,7 +18,6 @@
  */
 
 
-import org.apache.kafka.connect.data.Schema;
 import org.apache.kafka.connect.sink.SinkRecord;
 
 /**
@@ -27,10 +26,10 @@
  */
 public interface RecordConverter<R> {
   /**
-   * @param schema The schema of the record to convert.
-   * @param struct The struct of the record to convert.
+   * @param record The record to convert.
+   * @param recordType The type of the record to convert, either value or key.
    * @return The converted record.
    */
-  R convertRecord(Schema schema, Object struct);
+  R convertRecord(SinkRecord record, KafkaSchemaRecordType recordType);
 
 }

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

@@ -40,7 +40,7 @@ public void configure(Map<String, String> properties) {
   }
 
   @Override
-  public Schema retrieveSchema(TableId table, String topic, String subjectType) {
+  public Schema retrieveSchema(TableId table, String topic, String schemaType) {
     String tableName = table.getTable();
     Schema schema = schemaCache.get(getCacheKey(tableName, topic));
     if (schema != null) {