NO-SNOW: Expose the Ingest SDK MAX_CLIENT_LAG configuration in KC (#758)

pom.xml

@@ -334,7 +334,7 @@
         <dependency>
             <groupId>net.snowflake</groupId>
             <artifactId>snowflake-ingest-sdk</artifactId>
-            <version>2.0.3</version>
+            <version>2.0.4</version>
             <exclusions>
                 <exclusion>
                     <groupId>net.snowflake</groupId>

pom_confluent.xml

@@ -386,7 +386,7 @@
         <dependency>
             <groupId>net.snowflake</groupId>
             <artifactId>snowflake-ingest-sdk</artifactId>
-            <version>2.0.3</version>
+            <version>2.0.4</version>
             <exclusions>
                 <exclusion>
                     <groupId>net.snowflake</groupId>

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

@@ -117,9 +117,9 @@ public class SnowflakeSinkConnectorConfig {
   public static final String INGESTION_METHOD_DEFAULT_SNOWPIPE =
       IngestionMethodConfig.SNOWPIPE.toString();
 
-  // This is the streaming bdec file version which can be defined in config
-  // NOTE: Please do not override this value unless recommended from snowflake
-  public static final String SNOWPIPE_STREAMING_FILE_VERSION = "snowflake.streaming.file.version";
+  // This is the streaming max client lag which can be defined in config
+  public static final String SNOWPIPE_STREAMING_MAX_CLIENT_LAG =
+      "snowflake.streaming.max.client.lag";
 
   // TESTING
   public static final String REBALANCING = "snowflake.test.rebalancing";
@@ -541,17 +541,16 @@ static ConfigDef newConfigDef() {
             ConfigDef.Width.NONE,
             INGESTION_METHOD_OPT)
         .define(
-            SNOWPIPE_STREAMING_FILE_VERSION,
-            Type.STRING,
-            "", // default is handled in Ingest SDK
-            null, // no validator
+            SNOWPIPE_STREAMING_MAX_CLIENT_LAG,
+            Type.LONG,
+            StreamingUtils.STREAMING_BUFFER_FLUSH_TIME_MINIMUM_SEC,
+            ConfigDef.Range.atLeast(StreamingUtils.STREAMING_BUFFER_FLUSH_TIME_MINIMUM_SEC),
             Importance.LOW,
-            "Acceptable values for Snowpipe Streaming BDEC Versions: 1 and 3. Check Ingest"
-                + " SDK for default behavior. Please do not set this unless Absolutely needed. ",
+            "Decide how often the buffer in the Ingest SDK will be flushed",
             CONNECTOR_CONFIG,
             6,
             ConfigDef.Width.NONE,
-            SNOWPIPE_STREAMING_FILE_VERSION)
+            SNOWPIPE_STREAMING_MAX_CLIENT_LAG)
         .define(
             ERRORS_TOLERANCE_CONFIG,
             Type.STRING,

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

@@ -435,12 +435,12 @@ static String validateConfig(Map<String, String> config) {
                 "Schematization is only available with {}.",
                 IngestionMethodConfig.SNOWPIPE_STREAMING.toString()));
       }
-      if (config.containsKey(SnowflakeSinkConnectorConfig.SNOWPIPE_STREAMING_FILE_VERSION)) {
+      if (config.containsKey(SnowflakeSinkConnectorConfig.SNOWPIPE_STREAMING_MAX_CLIENT_LAG)) {
         invalidConfigParams.put(
-            SnowflakeSinkConnectorConfig.SNOWPIPE_STREAMING_FILE_VERSION,
+            SnowflakeSinkConnectorConfig.SNOWPIPE_STREAMING_MAX_CLIENT_LAG,
             Utils.formatString(
                 "{} is only available with ingestion type: {}.",
-                SnowflakeSinkConnectorConfig.SNOWPIPE_STREAMING_FILE_VERSION,
+                SnowflakeSinkConnectorConfig.SNOWPIPE_STREAMING_MAX_CLIENT_LAG,
                 IngestionMethodConfig.SNOWPIPE_STREAMING.toString()));
       }
       if (config.containsKey(

src/main/java/com/snowflake/kafka/connector/internal/streaming/StreamingClientProperties.java

@@ -17,14 +17,15 @@
 
 package com.snowflake.kafka.connector.internal.streaming;
 
-import static com.snowflake.kafka.connector.SnowflakeSinkConnectorConfig.SNOWPIPE_STREAMING_FILE_VERSION;
-import static net.snowflake.ingest.utils.ParameterProvider.BLOB_FORMAT_VERSION;
+import static com.snowflake.kafka.connector.SnowflakeSinkConnectorConfig.SNOWPIPE_STREAMING_MAX_CLIENT_LAG;
+import static net.snowflake.ingest.utils.ParameterProvider.MAX_CLIENT_LAG;
 
 import com.snowflake.kafka.connector.Utils;
 import com.snowflake.kafka.connector.internal.KCLogger;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.Properties;
 import java.util.stream.Collectors;
@@ -77,15 +78,15 @@ public StreamingClientProperties(Map<String, String> connectorConfig) {
         STREAMING_CLIENT_PREFIX_NAME
             + connectorConfig.getOrDefault(Utils.NAME, DEFAULT_CLIENT_NAME);
 
-    // Override only if bdec version is explicitly set in config, default to the version set
-    // inside Ingest SDK
+    // Override only if the max client lag is explicitly set in config
     this.parameterOverrides = new HashMap<>();
-    Optional<String> snowpipeStreamingBdecVersion =
-        Optional.ofNullable(connectorConfig.get(SNOWPIPE_STREAMING_FILE_VERSION));
-    snowpipeStreamingBdecVersion.ifPresent(
+    Optional<String> snowpipeStreamingMaxClientLag =
+        Optional.ofNullable(connectorConfig.get(SNOWPIPE_STREAMING_MAX_CLIENT_LAG));
+    snowpipeStreamingMaxClientLag.ifPresent(
         overriddenValue -> {
-          LOGGER.info("Config is overridden for {} ", SNOWPIPE_STREAMING_FILE_VERSION);
-          parameterOverrides.put(BLOB_FORMAT_VERSION, overriddenValue);
+          LOGGER.info(
+              "Config is overridden for {}={}", SNOWPIPE_STREAMING_MAX_CLIENT_LAG, overriddenValue);
+          parameterOverrides.put(MAX_CLIENT_LAG, String.format("%s second", overriddenValue));
         });
   }
 
@@ -118,7 +119,8 @@ public String getLoggableClientProperties() {
   @Override
   public boolean equals(Object other) {
     return other.getClass().equals(StreamingClientProperties.class)
-        & ((StreamingClientProperties) other).clientProperties.equals(this.clientProperties);
+        && ((StreamingClientProperties) other).clientProperties.equals(this.clientProperties)
+        && ((StreamingClientProperties) other).parameterOverrides.equals(this.parameterOverrides);
   }
 
   /**
@@ -129,6 +131,6 @@ public boolean equals(Object other) {
    */
   @Override
   public int hashCode() {
-    return this.clientProperties.hashCode();
+    return Objects.hash(this.clientProperties, this.parameterOverrides);
   }
 }

src/test/java/com/snowflake/kafka/connector/ConnectorConfigTest.java

@@ -746,12 +746,12 @@ public void testInvalidValueConvertersForStreamingSnowpipe() {
   public void testInValidConfigFileTypeForSnowpipe() {
     try {
       Map<String, String> config = getConfig();
-      config.put(SnowflakeSinkConnectorConfig.SNOWPIPE_STREAMING_FILE_VERSION, "3");
+      config.put(SnowflakeSinkConnectorConfig.SNOWPIPE_STREAMING_MAX_CLIENT_LAG, "3");
       Utils.validateConfig(config);
     } catch (SnowflakeKafkaConnectorException exception) {
       assert exception
           .getMessage()
-          .contains(SnowflakeSinkConnectorConfig.SNOWPIPE_STREAMING_FILE_VERSION);
+          .contains(SnowflakeSinkConnectorConfig.SNOWPIPE_STREAMING_MAX_CLIENT_LAG);
     }
   }
 
@@ -763,14 +763,14 @@ public void testValidFileTypesForSnowpipeStreaming() {
         IngestionMethodConfig.SNOWPIPE_STREAMING.toString());
     config.put(Utils.SF_ROLE, "ACCOUNTADMIN");
 
-    config.put(SnowflakeSinkConnectorConfig.SNOWPIPE_STREAMING_FILE_VERSION, "3");
+    config.put(SnowflakeSinkConnectorConfig.SNOWPIPE_STREAMING_MAX_CLIENT_LAG, "3");
     Utils.validateConfig(config);
 
-    config.put(SnowflakeSinkConnectorConfig.SNOWPIPE_STREAMING_FILE_VERSION, "1");
+    config.put(SnowflakeSinkConnectorConfig.SNOWPIPE_STREAMING_MAX_CLIENT_LAG, "1");
     Utils.validateConfig(config);
 
     // lower case
-    config.put(SnowflakeSinkConnectorConfig.SNOWPIPE_STREAMING_FILE_VERSION, "abcd");
+    config.put(SnowflakeSinkConnectorConfig.SNOWPIPE_STREAMING_MAX_CLIENT_LAG, "abcd");
     Utils.validateConfig(config);
   }
 

src/test/java/com/snowflake/kafka/connector/internal/streaming/SnowflakeSinkServiceV2IT.java

@@ -1,5 +1,6 @@
 package com.snowflake.kafka.connector.internal.streaming;
 
+import static com.snowflake.kafka.connector.SnowflakeSinkConnectorConfig.SNOWPIPE_STREAMING_MAX_CLIENT_LAG;
 import static com.snowflake.kafka.connector.internal.streaming.SnowflakeSinkServiceV2.partitionChannelKey;
 import static com.snowflake.kafka.connector.internal.streaming.TopicPartitionChannel.NO_OFFSET_TOKEN_REGISTERED_IN_SNOWFLAKE;
 
@@ -43,6 +44,7 @@
 import org.apache.kafka.connect.sink.SinkRecord;
 import org.junit.After;
 import org.junit.Assert;
+import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.Parameterized;
@@ -1461,12 +1463,51 @@ public void testSchematizationSchemaEvolutionWithNonNullableColumn() throws Exce
   }
 
   @Test
-  public void testStreamingIngestion_invalid_file_version() throws Exception {
+  @Ignore // SNOW-986359: disable for now due to failure in merge gate
+  public void testStreamingIngestionValidClientLag() throws Exception {
+    Map<String, String> config = getConfig();
+    config.put(SNOWPIPE_STREAMING_MAX_CLIENT_LAG, "30");
+    SnowflakeSinkConnectorConfig.setDefaultValues(config);
+    conn.createTable(table);
+
+    // opens a channel for partition 0, table and topic
+    SnowflakeSinkService service =
+        SnowflakeSinkServiceFactory.builder(conn, IngestionMethodConfig.SNOWPIPE_STREAMING, config)
+            .setRecordNumber(100)
+            .setFlushTime(1)
+            .setErrorReporter(new InMemoryKafkaRecordErrorReporter())
+            .setSinkTaskContext(new InMemorySinkTaskContext(Collections.singleton(topicPartition)))
+            .addTask(table, new TopicPartition(topic, partition)) // Internally calls startTask
+            .build();
+
+    final long noOfRecords = 50;
+    List<SinkRecord> sinkRecords =
+        TestUtils.createJsonStringSinkRecords(0, noOfRecords, topic, partition);
+
+    service.insert(sinkRecords);
+    try {
+      // Wait 20 seconds here and no flush should happen since the max client lag is 30 seconds
+      TestUtils.assertWithRetry(
+          () -> service.getOffset(new TopicPartition(topic, partition)) == noOfRecords, 5, 4);
+      Assert.fail("The rows should not be flushed");
+    } catch (Exception e) {
+      // do nothing
+    }
+
+    // Wait for enough time, the rows should be flushed
+    TestUtils.assertWithRetry(
+        () -> service.getOffset(new TopicPartition(topic, partition)) == noOfRecords, 30, 30);
+
+    service.closeAll();
+  }
+
+  @Test
+  public void testStreamingIngestionInvalidClientLag() {
     Map<String, String> config = TestUtils.getConfForStreaming();
     SnowflakeSinkConnectorConfig.setDefaultValues(config);
     Map<String, String> overriddenConfig = new HashMap<>(config);
     overriddenConfig.put(
-        SnowflakeSinkConnectorConfig.SNOWPIPE_STREAMING_FILE_VERSION, "TWOO_HUNDRED");
+        SnowflakeSinkConnectorConfig.SNOWPIPE_STREAMING_MAX_CLIENT_LAG, "TWOO_HUNDRED");
 
     conn.createTable(table);
 
@@ -1526,16 +1567,30 @@ public void testStreamingIngest_multipleChannel_distinctClients() throws Excepti
     dogConfig.put(Utils.SF_OAUTH_CLIENT_ID, "2");
     dogConfig.put(Utils.NAME, dogTopic);
 
+    String fishTopic = "fishTopic_" + TestUtils.randomTableName();
+    Map<String, String> fishConfig = getConfig();
+    SnowflakeSinkConnectorConfig.setDefaultValues(fishConfig);
+    fishConfig.put(
+        SnowflakeSinkConnectorConfig.ENABLE_STREAMING_CLIENT_OPTIMIZATION_CONFIG, "true");
+    fishConfig.put(Utils.SF_OAUTH_CLIENT_ID, "2");
+    fishConfig.put(Utils.NAME, fishTopic);
+    fishConfig.put(SNOWPIPE_STREAMING_MAX_CLIENT_LAG, "1");
+
     // setup connection and create tables
     TopicPartition catTp = new TopicPartition(catTopic, 0);
     SnowflakeConnectionService catConn =
         SnowflakeConnectionServiceFactory.builder().setProperties(catConfig).build();
     catConn.createTable(catTopic);
 
     TopicPartition dogTp = new TopicPartition(dogTopic, 1);
-    SnowflakeConnectionService dogconn =
+    SnowflakeConnectionService dogConn =
         SnowflakeConnectionServiceFactory.builder().setProperties(dogConfig).build();
-    dogconn.createTable(dogTopic);
+    dogConn.createTable(dogTopic);
+
+    TopicPartition fishTp = new TopicPartition(fishTopic, 1);
+    SnowflakeConnectionService fishConn =
+        SnowflakeConnectionServiceFactory.builder().setProperties(fishConfig).build();
+    fishConn.createTable(fishTopic);
 
     // create the sink services
     SnowflakeSinkService catService =
@@ -1549,48 +1604,70 @@ public void testStreamingIngest_multipleChannel_distinctClients() throws Excepti
 
     SnowflakeSinkService dogService =
         SnowflakeSinkServiceFactory.builder(
-                dogconn, IngestionMethodConfig.SNOWPIPE_STREAMING, dogConfig)
+                dogConn, IngestionMethodConfig.SNOWPIPE_STREAMING, dogConfig)
             .setRecordNumber(1)
             .setErrorReporter(new InMemoryKafkaRecordErrorReporter())
             .setSinkTaskContext(new InMemorySinkTaskContext(Collections.singleton(dogTp)))
             .addTask(dogTopic, dogTp) // Internally calls startTask
             .build();
 
+    SnowflakeSinkService fishService =
+        SnowflakeSinkServiceFactory.builder(
+                dogConn, IngestionMethodConfig.SNOWPIPE_STREAMING, fishConfig)
+            .setRecordNumber(1)
+            .setErrorReporter(new InMemoryKafkaRecordErrorReporter())
+            .setSinkTaskContext(new InMemorySinkTaskContext(Collections.singleton(fishTp)))
+            .addTask(fishTopic, fishTp) // Internally calls startTask
+            .build();
+
     // create records
     final int catRecordCount = 9;
     final int dogRecordCount = 3;
+    final int fishRecordCount = 1;
 
     List<SinkRecord> catRecords =
         TestUtils.createJsonStringSinkRecords(0, catRecordCount, catTp.topic(), catTp.partition());
     List<SinkRecord> dogRecords =
         TestUtils.createJsonStringSinkRecords(0, dogRecordCount, dogTp.topic(), dogTp.partition());
+    List<SinkRecord> fishRecords =
+        TestUtils.createJsonStringSinkRecords(
+            0, fishRecordCount, fishTp.topic(), fishTp.partition());
 
     // insert records
     catService.insert(catRecords);
     dogService.insert(dogRecords);
+    fishService.insert(fishRecords);
 
     // check data was ingested
     TestUtils.assertWithRetry(() -> catService.getOffset(catTp) == catRecordCount, 20, 20);
     TestUtils.assertWithRetry(() -> dogService.getOffset(dogTp) == dogRecordCount, 20, 20);
+    TestUtils.assertWithRetry(() -> fishService.getOffset(fishTp) == fishRecordCount, 20, 20);
 
-    // verify two clients were created
+    // verify three clients were created
     assert StreamingClientProvider.getStreamingClientProviderInstance()
         .getRegisteredClients()
         .containsKey(new StreamingClientProperties(catConfig));
     assert StreamingClientProvider.getStreamingClientProviderInstance()
         .getRegisteredClients()
         .containsKey(new StreamingClientProperties(dogConfig));
+    assert StreamingClientProvider.getStreamingClientProviderInstance()
+        .getRegisteredClients()
+        .containsKey(new StreamingClientProperties(fishConfig));
 
     // close services
     catService.closeAll();
     dogService.closeAll();
+    fishService.closeAll();
 
-    // verify both clients were closed
+    // verify three clients were closed
     assert !StreamingClientProvider.getStreamingClientProviderInstance()
         .getRegisteredClients()
         .containsKey(new StreamingClientProperties(catConfig));
     assert !StreamingClientProvider.getStreamingClientProviderInstance()
         .getRegisteredClients()
         .containsKey(new StreamingClientProperties(dogConfig));
+    assert !StreamingClientProvider.getStreamingClientProviderInstance()
+        .getRegisteredClients()
+        .containsKey(new StreamingClientProperties(fishConfig));
   }
 }

src/test/java/com/snowflake/kafka/connector/internal/streaming/StreamingClientHandlerTest.java

@@ -17,7 +17,6 @@
 
 package com.snowflake.kafka.connector.internal.streaming;
 
-import com.snowflake.kafka.connector.SnowflakeSinkConnectorConfig;
 import com.snowflake.kafka.connector.Utils;
 import com.snowflake.kafka.connector.internal.TestUtils;
 import java.util.Map;
@@ -87,15 +86,6 @@ public void testCreateClientException() {
     }
   }
 
-  @Test(expected = IllegalArgumentException.class)
-  public void testCreateClientInvalidBdecVersion() {
-    // add invalid bdec version
-    this.connectorConfig.put(SnowflakeSinkConnectorConfig.SNOWPIPE_STREAMING_FILE_VERSION, "1");
-
-    // test create
-    this.streamingClientHandler.createClient(new StreamingClientProperties(this.connectorConfig));
-  }
-
   @Test
   public void testCloseClient() throws Exception {
     // setup valid client