SNOW-121816 added additional flags to control the amount of metadata

changed to use snowflake metadata config object to pass config

added more assert to config testting

typo

changed function call to constructor

call construc from another construc

Author: sfc-gh-zli

src/main/java/com/snowflake/kafka/connector/SnowflakeSinkConnectorConfig.java

@@ -43,12 +43,12 @@ public class SnowflakeSinkConnectorConfig
   public static final long BUFFER_SIZE_BYTES_MAX = 100000000;
   static final String TOPICS_TABLES_MAP = "snowflake.topic2table.map";
 
-  //in second
+  // Time in seconds
   public static final long BUFFER_FLUSH_TIME_SEC_MIN = 10;
   public static final long BUFFER_FLUSH_TIME_SEC_DEFAULT = 30;
   static final String BUFFER_FLUSH_TIME_SEC = "buffer.flush.time";
 
-  // snowflake connection and database config
+  // Snowflake connection and database config
   private static final String SNOWFLAKE_LOGIN_INFO = "Snowflake Login Info";
   static final String SNOWFLAKE_URL = Utils.SF_URL;
   static final String SNOWFLAKE_USER = Utils.SF_USER;
@@ -57,19 +57,29 @@ public class SnowflakeSinkConnectorConfig
   static final String SNOWFLAKE_SCHEMA = Utils.SF_SCHEMA;
   static final String SNOWFLAKE_PRIVATE_KEY_PASSPHRASE = Utils.PRIVATE_KEY_PASSPHRASE;
 
-  //Proxy Info
+  // Proxy Info
   private static final String PROXY_INFO = "Proxy Info";
   static final String JVM_PROXY_HOST = "jvm.proxy.host";
   static final String JVM_PROXY_PORT = "jvm.proxy.port";
 
-  //Schema Registry Info
+  // Schema Registry Info
   private static final String SCHEMA_REGISTRY_INFO = "Schema Registry Info";
   static final String SCHEMA_REGISTRY_AUTH_CREDENTIALS_SOURCE =
     "value.converter.basic.auth.credentials.source";
   static final String SCHEMA_REGISTRY_AUTH_USER_INFO =
     "value.converter.basic.auth.user.info";
   private static final String REGISTRY_URL = "value.converter.schema.registry.url";
 
+  // Snowflake Metadata Flags
+  private static final String SNOWFLAKE_METADATA_FLAGS = "Snowflake Metadata Flags";
+  public static final String SNOWFLAKE_METADATA_CREATETIME = "snowflake.metadata.createtime";
+  public static final String SNOWFLAKE_METADATA_TOPIC = "snowflake.metadata.topic";
+  public static final String SNOWFLAKE_METADATA_OFFSET_AND_PARTITION =
+    "snowflake.metadata.offset.and.partition";
+  public static final String SNOWFLAKE_METADATA_ALL = "snowflake.metadata.all";
+  public static final String SNOWFLAKE_METADATA_DEFAULT = "true";
+
+
   static ConfigDef newConfigDef()
   {
     return new ConfigDef()
@@ -214,6 +224,42 @@ static ConfigDef newConfigDef()
         3,
         ConfigDef.Width.NONE,
         BUFFER_FLUSH_TIME_SEC)
+      .define(SNOWFLAKE_METADATA_ALL,
+        Type.STRING,
+        SNOWFLAKE_METADATA_DEFAULT,
+        Importance.LOW,
+        "Flag to control whether there is metadata collected. If set to false, all metadata will be dropped",
+        SNOWFLAKE_METADATA_FLAGS,
+        0,
+        ConfigDef.Width.NONE,
+        SNOWFLAKE_METADATA_ALL)
+      .define(SNOWFLAKE_METADATA_CREATETIME,
+        Type.STRING,
+        SNOWFLAKE_METADATA_DEFAULT,
+        Importance.LOW,
+        "Flag to control whether createtime is collected in snowflake metadata",
+        SNOWFLAKE_METADATA_FLAGS,
+        1,
+        ConfigDef.Width.NONE,
+        SNOWFLAKE_METADATA_CREATETIME)
+      .define(SNOWFLAKE_METADATA_TOPIC,
+        Type.STRING,
+        SNOWFLAKE_METADATA_DEFAULT,
+        Importance.LOW,
+        "Flag to control whether kafka topic name is collected in snowflake metadata",
+        SNOWFLAKE_METADATA_FLAGS,
+        2,
+        ConfigDef.Width.NONE,
+        SNOWFLAKE_METADATA_TOPIC)
+      .define(SNOWFLAKE_METADATA_OFFSET_AND_PARTITION,
+        Type.STRING,
+        SNOWFLAKE_METADATA_DEFAULT,
+        Importance.LOW,
+        "Flag to control whether kafka partition and offset are collected in snowflake metadata",
+        SNOWFLAKE_METADATA_FLAGS,
+        3,
+        ConfigDef.Width.NONE,
+        SNOWFLAKE_METADATA_OFFSET_AND_PARTITION)
       ;
   }
 }

src/main/java/com/snowflake/kafka/connector/SnowflakeSinkTask.java

@@ -22,6 +22,7 @@
 import com.snowflake.kafka.connector.internal.SnowflakeErrors;
 import com.snowflake.kafka.connector.internal.SnowflakeSinkService;
 import com.snowflake.kafka.connector.internal.SnowflakeSinkServiceFactory;
+import com.snowflake.kafka.connector.records.SnowflakeMetadataConfig;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.connect.errors.RetriableException;
@@ -108,6 +109,9 @@ public void start(final Map<String, String> parsedConfig)
     //generate topic to table map
     this.topic2table = getTopicToTableMap(parsedConfig);
 
+    // generate metadataConfig table
+    SnowflakeMetadataConfig metadataConfig = new SnowflakeMetadataConfig(parsedConfig);
+
     //enable jvm proxy
     Utils.enableJVMProxy(parsedConfig);
 
@@ -135,6 +139,7 @@ public void start(final Map<String, String> parsedConfig)
       .setRecordNumber(bufferCountRecords)
       .setFlushTime(bufferFlushTime)
       .setTopic2TableMap(topic2table)
+      .setMetadataConfig(metadataConfig)
       .build();
   }
 

src/main/java/com/snowflake/kafka/connector/internal/SnowflakeSinkService.java

@@ -1,5 +1,6 @@
 package com.snowflake.kafka.connector.internal;
 
+import com.snowflake.kafka.connector.records.SnowflakeMetadataConfig;
 import net.snowflake.client.jdbc.internal.fasterxml.jackson.databind.JsonNode;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.connect.sink.SinkRecord;
@@ -81,6 +82,12 @@ public interface SnowflakeSinkService
    */
   void setFlushTime(long time);
 
+  /**
+   * set the metadata config to let user control what metadata to be collected into SF db
+   * @param configMap a String to String Map
+   */
+  void setMetadataConfig(SnowflakeMetadataConfig configMap);
+
   /**
    * @return current number of record limitation
    */

src/main/java/com/snowflake/kafka/connector/internal/SnowflakeSinkServiceFactory.java

@@ -1,5 +1,7 @@
 package com.snowflake.kafka.connector.internal;
 
+import com.snowflake.kafka.connector.records.SnowflakeMetadataConfig;
+
 import java.security.PrivateKey;
 import java.util.Map;
 
@@ -62,6 +64,7 @@ public SnowflakeSinkServiceBuilder setFlushTime(long time)
       logInfo("flush time is limited to {}", time);
       return this;
     }
+
     public SnowflakeSinkServiceBuilder setTopic2TableMap(Map<String, String> topic2TableMap)
     {
       this.service.setTopic2TableMap(topic2TableMap);
@@ -74,6 +77,13 @@ public SnowflakeSinkServiceBuilder setTopic2TableMap(Map<String, String> topic2T
       return this;
     }
 
+    public SnowflakeSinkServiceBuilder setMetadataConfig(SnowflakeMetadataConfig configMap)
+    {
+      this.service.setMetadataConfig(configMap);
+      logInfo("metadata config map is {}", configMap.toString());
+      return this;
+    }
+
     public SnowflakeSinkService build()
     {
       logInfo("{} created", SnowflakeSinkService.class.getName());

src/main/java/com/snowflake/kafka/connector/internal/SnowflakeSinkServiceV1.java

@@ -3,6 +3,7 @@
 import com.snowflake.kafka.connector.SnowflakeSinkConnectorConfig;
 import com.snowflake.kafka.connector.Utils;
 import com.snowflake.kafka.connector.records.RecordService;
+import com.snowflake.kafka.connector.records.SnowflakeMetadataConfig;
 import com.snowflake.kafka.connector.records.SnowflakeRecordContent;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.connect.sink.SinkRecord;
@@ -206,6 +207,12 @@ public void setTopic2TableMap(Map<String, String> topic2TableMap)
     this.topic2TableMap = topic2TableMap;
   }
 
+  @Override
+  public void setMetadataConfig(SnowflakeMetadataConfig configMap)
+  {
+    this.recordService.setMetadataConfig(configMap);
+  }
+
   @Override
   public long getRecordNumber()
   {

src/main/java/com/snowflake/kafka/connector/records/RecordService.java

@@ -16,6 +16,7 @@
  */
 package com.snowflake.kafka.connector.records;
 
+import com.snowflake.kafka.connector.SnowflakeSinkConnectorConfig;
 import com.snowflake.kafka.connector.internal.Logging;
 import com.snowflake.kafka.connector.internal.SnowflakeErrors;
 import net.snowflake.client.jdbc.internal.fasterxml.jackson.databind.JsonNode;
@@ -32,21 +33,26 @@
 import org.apache.kafka.connect.header.Headers;
 import org.apache.kafka.connect.sink.SinkRecord;
 
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
 public class RecordService extends Logging
 {
   private static final ObjectMapper MAPPER = new ObjectMapper();
 
-  private static final String OFFSET = "offset";
-  private static final String TOPIC = "topic";
-  private static final String PARTITION = "partition";
-  private static final String KEY = "key";
-  private static final String CONTENT = "content";
-  private static final String META = "meta";
-  private static final String SCHEMA_ID = "schema_id";
-  private static final String HEADERS = "headers";
+  // deleted private to use these values in test
+  static final String OFFSET = "offset";
+  static final String TOPIC = "topic";
+  static final String PARTITION = "partition";
+  static final String KEY = "key";
+  static final String CONTENT = "content";
+  static final String META = "meta";
+  static final String SCHEMA_ID = "schema_id";
+  static final String HEADERS = "headers";
+
+  // This class is designed to work with empty metadata config map
+  private SnowflakeMetadataConfig metadataConfig = new SnowflakeMetadataConfig();
 
   /**
    * process records
@@ -68,6 +74,9 @@ public RecordService()
   {
   }
 
+  public void setMetadataConfig(SnowflakeMetadataConfig metadataConfigIn) {
+    metadataConfig = metadataConfigIn;
+  }
 
   /**
    * process given SinkRecord,
@@ -91,12 +100,19 @@ public String processRecord(SinkRecord record)
     SnowflakeRecordContent content = (SnowflakeRecordContent) record.value();
 
     ObjectNode meta = MAPPER.createObjectNode();
-    meta.put(OFFSET, record.kafkaOffset());
-    meta.put(TOPIC, record.topic());
-    meta.put(PARTITION, record.kafkaPartition());
+    if (metadataConfig.topicFlag)
+    {
+      meta.put(TOPIC, record.topic());
+    }
+    if (metadataConfig.offsetAndPartitionFlag)
+    {
+      meta.put(OFFSET, record.kafkaOffset());
+      meta.put(PARTITION, record.kafkaPartition());
+    }
 
     //ignore if no timestamp
-    if (record.timestampType() != TimestampType.NO_TIMESTAMP_TYPE)
+    if (record.timestampType() != TimestampType.NO_TIMESTAMP_TYPE &&
+      metadataConfig.createtimeFlag)
     {
       meta.put(record.timestampType().name, record.timestamp());
     }
@@ -125,7 +141,10 @@ public String processRecord(SinkRecord record)
     {
       ObjectNode data = MAPPER.createObjectNode();
       data.set(CONTENT, node);
-      data.set(META, meta);
+      if (metadataConfig.allFlag)
+      {
+        data.set(META, meta);
+      }
       buffer.append(data.toString());
     }
     return buffer.toString();

src/main/java/com/snowflake/kafka/connector/records/SnowflakeMetadataConfig.java

@@ -0,0 +1,74 @@
+package com.snowflake.kafka.connector.records;
+
+import com.snowflake.kafka.connector.SnowflakeSinkConnectorConfig;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class SnowflakeMetadataConfig {
+  final boolean createtimeFlag;
+  final boolean topicFlag;
+  final boolean offsetAndPartitionFlag;
+  final boolean allFlag;
+
+  /**
+   * initialize with default config
+   */
+  public SnowflakeMetadataConfig()
+  {
+    this(new HashMap<String, String>());
+  }
+
+  /**
+   * set flag to false only if metadata config is set to false in config
+   * @param config a String to String map of configs
+   */
+  public SnowflakeMetadataConfig(Map<String, String> config)
+  {
+    // have those local variable to avoid assigning to final values multiple times
+    // these values are the default values of the configuration
+    boolean createtime = true;
+    boolean topic = true;
+    boolean offsetAndPartition = true;
+    boolean all = true;
+    if (config.containsKey(SnowflakeSinkConnectorConfig.SNOWFLAKE_METADATA_CREATETIME) &&
+        !config.get(SnowflakeSinkConnectorConfig.SNOWFLAKE_METADATA_CREATETIME)
+          .equals(SnowflakeSinkConnectorConfig.SNOWFLAKE_METADATA_DEFAULT))
+    {
+      createtime = false;
+    }
+    if (config.containsKey(SnowflakeSinkConnectorConfig.SNOWFLAKE_METADATA_TOPIC) &&
+      !config.get(SnowflakeSinkConnectorConfig.SNOWFLAKE_METADATA_TOPIC)
+        .equals(SnowflakeSinkConnectorConfig.SNOWFLAKE_METADATA_DEFAULT))
+    {
+      topic = false;
+    }
+    if (config.containsKey(SnowflakeSinkConnectorConfig.SNOWFLAKE_METADATA_OFFSET_AND_PARTITION) &&
+      !config.get(SnowflakeSinkConnectorConfig.SNOWFLAKE_METADATA_OFFSET_AND_PARTITION)
+        .equals(SnowflakeSinkConnectorConfig.SNOWFLAKE_METADATA_DEFAULT))
+    {
+      offsetAndPartition = false;
+    }
+    if (config.containsKey(SnowflakeSinkConnectorConfig.SNOWFLAKE_METADATA_ALL) &&
+      !config.get(SnowflakeSinkConnectorConfig.SNOWFLAKE_METADATA_ALL)
+        .equals(SnowflakeSinkConnectorConfig.SNOWFLAKE_METADATA_DEFAULT))
+    {
+      all = false;
+    }
+
+    createtimeFlag = createtime;
+    topicFlag = topic;
+    offsetAndPartitionFlag = offsetAndPartition;
+    allFlag = all;
+
+  }
+
+  public String toString()
+  {
+    return
+      "{createtimeFlag: " + createtimeFlag + ", " +
+      "topicFlag: " + topicFlag + ", " +
+      "offsetAndPartitionFlag: " + offsetAndPartitionFlag + ", " +
+      "allFlag: " + allFlag + "}";
+  }
+}