SNOW-1623269 Fail sink task on authorization exception from Snowflake (…

…#916)

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

@@ -214,6 +214,10 @@ public class SnowflakeSinkConnectorConfig {
  + " format is deprecated and V1 will be used always, disabling this config could have"
  + " ramifications. Please consult Snowflake support before setting this to false.";
 
+ public static final String ENABLE_TASK_FAIL_ON_AUTHORIZATION_ERRORS =
+ "enable.task.fail.on.authorization.errors";
+ public static final boolean ENABLE_TASK_FAIL_ON_AUTHORIZATION_ERRORS_DEFAULT = false;
+
  // MDC logging header
  public static final String ENABLE_MDC_LOGGING_CONFIG = "enable.mdc.logging";
  public static final String ENABLE_MDC_LOGGING_DISPLAY = "Enable MDC logging";
@@ -714,7 +718,14 @@ static ConfigDef newConfigDef() {
  CONNECTOR_CONFIG,
  9,
  ConfigDef.Width.NONE,
- ENABLE_CHANNEL_OFFSET_TOKEN_MIGRATION_DISPLAY);
+ ENABLE_CHANNEL_OFFSET_TOKEN_MIGRATION_DISPLAY)
+ .define(
+ ENABLE_TASK_FAIL_ON_AUTHORIZATION_ERRORS,
+ Type.BOOLEAN,
+ ENABLE_TASK_FAIL_ON_AUTHORIZATION_ERRORS_DEFAULT,
+ Importance.LOW,
+ "If set to true the Connector will fail its tasks when authorization error from"
+ + " Snowflake occurred");
  }
 
  public static class TopicToTableValidator implements ConfigDef.Validator {

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

@@ -89,6 +89,9 @@ public class SnowflakeSinkTask extends SinkTask {
 
  private IngestionMethodConfig ingestionMethodConfig;
 
+ private final SnowflakeSinkTaskAuthorizationExceptionTracker authorizationExceptionTracker =
+ new SnowflakeSinkTaskAuthorizationExceptionTracker();
+
  /** default constructor, invoked by kafka connect framework */
  public SnowflakeSinkTask() {
  DYNAMIC_LOGGER = new KCLogger(this.getClass().getName());
@@ -156,6 +159,8 @@ public void start(final Map<String, String> parsedConfig) {
  // generate topic to table map
  this.topic2table = getTopicToTableMap(parsedConfig);
 
+ this.authorizationExceptionTracker.updateStateOnTaskStart(parsedConfig);
+
  // generate metadataConfig table
  SnowflakeMetadataConfig metadataConfig = new SnowflakeMetadataConfig(parsedConfig);
 
@@ -294,6 +299,8 @@ public void close(final Collection<TopicPartition> partitions) {
  */
  @Override
  public void put(final Collection<SinkRecord> records) {
+ this.authorizationExceptionTracker.throwExceptionIfAuthorizationFailed();
+
  final long recordSize = records.size();
  if (enableRebalancing && recordSize > 0) {
  processRebalancingTest();
@@ -345,6 +352,7 @@ public Map<TopicPartition, OffsetAndMetadata> preCommit(
  }
  });
  } catch (Exception e) {
+ this.authorizationExceptionTracker.reportPrecommitException(e);
  this.DYNAMIC_LOGGER.error("PreCommit error: {} ", e.getMessage());
  }
 

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

@@ -0,0 +1,54 @@
+package com.snowflake.kafka.connector;
+
+import static com.snowflake.kafka.connector.SnowflakeSinkConnectorConfig.ENABLE_TASK_FAIL_ON_AUTHORIZATION_ERRORS;
+import static com.snowflake.kafka.connector.SnowflakeSinkConnectorConfig.ENABLE_TASK_FAIL_ON_AUTHORIZATION_ERRORS_DEFAULT;
+import static com.snowflake.kafka.connector.internal.SnowflakeErrors.ERROR_1005;
+
+import java.util.Map;
+
+/**
+ * When the user rotates Snowflake key that is stored in an external file the Connector hangs and
+ * does not mark its tasks as failed. To fix this corner case we need to track the authorization
+ * exception thrown during preCommit() and stop tasks during put().
+ *
+ * <p>Note that exceptions thrown during preCommit() are swallowed by Kafka Connect and will not
+ * cause task failure.
+ */
+public class SnowflakeSinkTaskAuthorizationExceptionTracker {
+
+ private static final String AUTHORIZATION_EXCEPTION_MESSAGE = "Authorization failed after retry";
+
+ private boolean authorizationTaskFailureEnabled;
+ private boolean authorizationErrorReported;
+
+ public SnowflakeSinkTaskAuthorizationExceptionTracker() {
+ this.authorizationTaskFailureEnabled = true;
+ this.authorizationErrorReported = false;
+ }
+
+ public void updateStateOnTaskStart(Map<String, String> taskConfig) {
+ authorizationTaskFailureEnabled =
+ Boolean.parseBoolean(
+ taskConfig.getOrDefault(
+ ENABLE_TASK_FAIL_ON_AUTHORIZATION_ERRORS,
+ Boolean.toString(ENABLE_TASK_FAIL_ON_AUTHORIZATION_ERRORS_DEFAULT)));
+ }
+
+ /**
+ * Check if the thrown exception is related to authorization
+ *
+ * @param ex - any exception that occurred during preCommit
+ */
+ public void reportPrecommitException(Exception ex) {
+ if (ex.getMessage().contains(AUTHORIZATION_EXCEPTION_MESSAGE)) {
+ authorizationErrorReported = true;
+ }
+ }
+
+ /** Throw exception if authorization has failed before */
+ public void throwExceptionIfAuthorizationFailed() {
+ if (authorizationTaskFailureEnabled && authorizationErrorReported) {
+ throw ERROR_1005.getException();
+ }
+ }
+}

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

@@ -155,6 +155,10 @@ public enum SnowflakeErrors {
  "Either the current connection is closed or hasn't connect to snowflake" + " server"),
  ERROR_1004(
  "1004", "Fetching OAuth token fail", "Fail to get OAuth token from authorization server"),
+ ERROR_1005(
+ "1005",
+ "Task failed due to authorization error",
+ "Set `enable.task.fail.on.authorization.errors=false` to avoid this behavior"),
  // SQL issues 2---
  ERROR_2001(
  "2001", "Failed to prepare SQL statement", "SQL Exception, reported by Snowflake JDBC"),

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

@@ -0,0 +1,69 @@
+package com.snowflake.kafka.connector;
+
+import static com.snowflake.kafka.connector.SnowflakeSinkConnectorConfig.ENABLE_TASK_FAIL_ON_AUTHORIZATION_ERRORS;
+
+import com.snowflake.kafka.connector.internal.SnowflakeKafkaConnectorException;
+import com.snowflake.kafka.connector.internal.TestUtils;
+import java.util.Map;
+import java.util.stream.Stream;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+
+class SnowflakeSinkTaskAuthorizationExceptionTrackerTest {
+
+ @Test
+ public void shouldThrowExceptionOnAuthorizationError() {
+ // given
+ SnowflakeSinkTaskAuthorizationExceptionTracker tracker =
+ new SnowflakeSinkTaskAuthorizationExceptionTracker();
+ Map<String, String> config = TestUtils.getConfig();
+ config.put(ENABLE_TASK_FAIL_ON_AUTHORIZATION_ERRORS, "true");
+ tracker.updateStateOnTaskStart(config);
+
+ // when
+ tracker.reportPrecommitException(new Exception("Authorization failed after retry"));
+
+ // then
+ Assertions.assertThrows(
+ SnowflakeKafkaConnectorException.class, tracker::throwExceptionIfAuthorizationFailed);
+ }
+
+ @Test
+ public void shouldNotThrowExceptionWhenNoExceptionReported() {
+ // given
+ SnowflakeSinkTaskAuthorizationExceptionTracker tracker =
+ new SnowflakeSinkTaskAuthorizationExceptionTracker();
+ Map<String, String> config = TestUtils.getConfig();
+ config.put(ENABLE_TASK_FAIL_ON_AUTHORIZATION_ERRORS, "true");
+ tracker.updateStateOnTaskStart(config);
+
+ // expect
+ Assertions.assertDoesNotThrow(tracker::throwExceptionIfAuthorizationFailed);
+ }
+
+ @ParameterizedTest
+ @MethodSource("noExceptionConditions")
+ public void shouldNotThrowException(boolean enabled, String exceptionMessage) {
+ // given
+ SnowflakeSinkTaskAuthorizationExceptionTracker tracker =
+ new SnowflakeSinkTaskAuthorizationExceptionTracker();
+ Map<String, String> config = TestUtils.getConfig();
+ config.put(ENABLE_TASK_FAIL_ON_AUTHORIZATION_ERRORS, Boolean.toString(enabled));
+ tracker.updateStateOnTaskStart(config);
+
+ // when
+ tracker.reportPrecommitException(new Exception(exceptionMessage));
+
+ // then
+ Assertions.assertDoesNotThrow(tracker::throwExceptionIfAuthorizationFailed);
+ }
+
+ public static Stream<Arguments> noExceptionConditions() {
+ return Stream.of(
+ Arguments.of(false, "Authorization failed after retry"),
+ Arguments.of(true, "NullPointerException"));
+ }
+}