CONN-10511 Automatic failover for Snowpipe Streaming (#1238)

Co-authored-by: Lukasz Kucharski <[email protected]>

Author: sfc-gh-mbobowski

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

@@ -349,4 +349,12 @@ ChannelMigrateOffsetTokenResponseDTO migrateStreamingChannelOffsetToken(
    *     otherwise.
    */
   Optional<List<DescribeTableRow>> describeTable(String tableName);
+
+  /**
+   * Checks if the underlying JDBC connection is valid.
+   *
+   * @param timeoutSeconds the timeout in seconds to wait for the validation operation to complete
+   * @return true if the connection is valid, false otherwise
+   */
+  boolean isValid(int timeoutSeconds);
 }

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

@@ -1221,6 +1221,16 @@ protected ChannelMigrateOffsetTokenResponseDTO getChannelMigrateOffsetTokenRespo
     return channelMigrateOffsetTokenResponseDTO;
   }
 
+  @Override
+  public boolean isValid(int timeoutSeconds) {
+    try {
+      return this.conn != null && this.conn.isValid(timeoutSeconds);
+    } catch (SQLException e) {
+      LOGGER.warn("Error checking connection validity: {}", e.getMessage());
+      return false;
+    }
+  }
+
   public static class FormattingUtils {
     /**
      * Transform the objectName to uppercase unless it is enclosed in double quotes

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

@@ -13,13 +13,21 @@
  * backoff and jitter.
  *
  * <p>This class provides a clean interface to execute channel opening operations with automatic
- * retry on HTTP 429 (rate limiting) errors from Snowflake streaming service.
+ * retry on:
+ *
+ * <ul>
+ *   <li>HTTP 429 (rate limiting) errors from Snowflake streaming service
+ *   <li>Status code 47 (secondary database) errors during failover scenarios
+ * </ul>
  */
 class OpenChannelRetryPolicy {
 
   private static final KCLogger LOGGER = new KCLogger(OpenChannelRetryPolicy.class.getName());
 
   private static final String RATE_LIMIT_MESSAGE_PART = "HTTP Status: 429";
+  private static final String SECONDARY_DATABASE_MESSAGE_PART =
+      "Cannot ingest data into a table that is part of a secondary database";
+  private static final String SECONDARY_DATABASE_STATUS_CODE = "\"status_code\" : 47";
 
   // Retry policy constants
   /** Initial delay before the first retry attempt. */
@@ -37,8 +45,14 @@ class OpenChannelRetryPolicy {
   /**
    * Executes the provided channel opening action with retry handling.
    *
-   * <p>On SFException containing "429" (HTTP rate limiting), it will retry with exponential backoff
-   * and jitter with unlimited retry attempts. Other exceptions are not retried.
+   * <p>Retries with exponential backoff and jitter with unlimited retry attempts on:
+   *
+   * <ul>
+   *   <li>SFException containing "HTTP Status: 429" (rate limiting)
+   *   <li>SFException containing status code 47 (secondary database during failover)
+   * </ul>
+   *
+   * Other exceptions are not retried.
    *
    * @param channelOpenAction the action to execute (typically openChannelForTable call)
    * @param channelName the channel name for logging purposes
@@ -66,6 +80,20 @@ static SnowflakeStreamingIngestChannel executeWithRetry(
   }
 
   private static boolean isRetryableError(Throwable e) {
-    return e instanceof SFException && e.getMessage().contains(RATE_LIMIT_MESSAGE_PART);
+    if (!(e instanceof SFException)) {
+      return false;
+    }
+
+    String message = e.getMessage();
+
+    // Retry on HTTP 429 (rate limiting)
+    boolean isRateLimitingError = message.contains(RATE_LIMIT_MESSAGE_PART);
+    // Retry on status code 47 (secondary database during failover)
+    //
+    boolean isFailoverError =
+        message.contains(SECONDARY_DATABASE_STATUS_CODE)
+            || message.contains(SECONDARY_DATABASE_MESSAGE_PART);
+
+    return isRateLimitingError || isFailoverError;
   }
 }

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

@@ -14,6 +14,7 @@
 import com.snowflake.kafka.connector.dlq.KafkaRecordErrorReporter;
 import com.snowflake.kafka.connector.internal.KCLogger;
 import com.snowflake.kafka.connector.internal.SnowflakeConnectionService;
+import com.snowflake.kafka.connector.internal.SnowflakeConnectionServiceFactory;
 import com.snowflake.kafka.connector.internal.SnowflakeErrors;
 import com.snowflake.kafka.connector.internal.SnowflakeSinkService;
 import com.snowflake.kafka.connector.internal.metrics.MetricsJmxReporter;
@@ -60,7 +61,7 @@ public class SnowflakeSinkServiceV2 implements SnowflakeSinkService {
   private static final KCLogger LOGGER = new KCLogger(SnowflakeSinkServiceV2.class.getName());
 
   // Used to connect to Snowflake, could be null during testing
-  private final SnowflakeConnectionService conn;
+  private volatile SnowflakeConnectionService conn;
 
   private final RecordService recordService;
   private final SnowflakeTelemetryService telemetryService;
@@ -350,6 +351,11 @@ public void insert(SinkRecord record) {
           "Topic: {} Partition: {} hasn't been initialized by OPEN function",
           record.topic(),
           record.kafkaPartition());
+
+      // Check connection validity and recreate if needed before starting partition
+      // Needed to handle failover scenario
+      recreateInvalidConnection();
+
       startPartition(
           Utils.tableName(record.topic(), this.topicToTableMap),
           new TopicPartition(record.topic(), record.kafkaPartition()));
@@ -714,4 +720,30 @@ private void populateSchemaEvolutionPermissions(String tableName) {
       }
     }
   }
+
+  private void recreateInvalidConnection() {
+    try {
+      // Check if connection is null, closed, or invalid
+      boolean shouldRecreate = false;
+      if (conn == null || conn.isClosed()) {
+        shouldRecreate = true;
+      } else if (!conn.isValid(5)) {
+        shouldRecreate = true;
+        try {
+          conn.close();
+        } catch (Exception e) {
+          LOGGER.warn("Could not close the old connection before opening the new one.", e);
+        }
+      }
+      if (shouldRecreate) {
+        LOGGER.warn("Connection is invalid, attempting to recreate");
+        this.conn =
+            SnowflakeConnectionServiceFactory.builder().setProperties(connectorConfig).build();
+
+        LOGGER.info("Successfully recreated Snowflake connection");
+      }
+    } catch (Exception e) {
+      LOGGER.error("Failed to recreate connection: {}", e.getMessage());
+    }
+  }
 }

test/connect-log4j.properties

@@ -43,5 +43,5 @@ log4j.logger.org.apache.zookeeper=ERROR
 log4j.logger.org.reflections=ERROR
 
 # enable debug level logging for SF KC and Ingest SDK
-log4j.logger.com.snowflake.kafka.connector=DEBUG
-log4j.logger.net.snowflake=DEBUG
+log4j.logger.com.snowflake.kafka.connector=INFO
+log4j.logger.net.snowflake=INFO