Skip to content

Commit 3beb03e

Browse files
sfc-gh-lkucharskisfc-gh-lkucharski
committed
Fix the tests
1 parent e2d9666 commit 3beb03e

File tree

7 files changed

+281
-6
lines changed

7 files changed

+281
-6
lines changed

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -233,7 +233,11 @@ public enum SnowflakeErrors {
233233
"5029",
234234
"Destination table does not exist",
235235
"Destination table does not exist. Please ensure the destination table exists in Snowflake"
236-
+ " before starting the connector.");
236+
+ " before starting the connector."),
237+
ERROR_5030(
238+
"5030",
239+
"Channel error count threshold exceeded",
240+
"Channel has reported errors during data ingestion. Check the channel history for details.");
237241

238242
// properties
239243

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

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,8 @@ public class SnowflakeSinkServiceV2 implements SnowflakeSinkService {
8383
// default is true unless the configuration provided is false;
8484
// If this is true, we will enable Mbean for required classes and emit JMX metrics for monitoring
8585
private boolean enableCustomJMXMonitoring = KafkaConnectorConfigParams.JMX_OPT_DEFAULT;
86+
// Whether to tolerate errors during ingestion (based on errors.tolerance config)
87+
private final boolean tolerateErrors;
8688

8789
public SnowflakeSinkServiceV2(
8890
SnowflakeConnectionService conn,
@@ -121,11 +123,13 @@ public SnowflakeSinkServiceV2(
121123
}
122124

123125
this.metricsJmxReporter = new MetricsJmxReporter(new MetricRegistry(), this.connectorName);
126+
this.tolerateErrors = StreamingUtils.tolerateErrors(connectorConfig);
124127

125128
LOGGER.info(
126-
"SnowflakeSinkServiceV2 initialized for connector: {}, task: {}",
129+
"SnowflakeSinkServiceV2 initialized for connector: {}, task: {}, tolerateErrors: {}",
127130
this.connectorName,
128-
this.taskId);
131+
this.taskId,
132+
this.tolerateErrors);
129133
}
130134

131135
/** Gets a unique identifier consisting of connector name, topic name and partition number. */
@@ -325,8 +329,10 @@ public long getOffset(TopicPartition topicPartition) {
325329
String partitionChannelKey =
326330
makeChannelName(this.connectorName, topicPartition.topic(), topicPartition.partition());
327331
if (partitionsToChannel.containsKey(partitionChannelKey)) {
328-
long offset = partitionsToChannel.get(partitionChannelKey).getOffsetSafeToCommitToKafka();
329-
partitionsToChannel.get(partitionChannelKey).setLatestConsumerGroupOffset(offset);
332+
TopicPartitionChannel channel = partitionsToChannel.get(partitionChannelKey);
333+
channel.checkChannelStatusAndLogErrors(tolerateErrors);
334+
long offset = channel.getOffsetSafeToCommitToKafka();
335+
channel.setLatestConsumerGroupOffset(offset);
330336
LOGGER.info(
331337
"Fetched snowflake commited offset: [{}] for channel [{}]", offset, partitionChannelKey);
332338
return offset;

‎src/main/java/com/snowflake/kafka/connector/internal/streaming/channel/TopicPartitionChannel.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,8 @@ default long getOffsetSafeToCommitToKafka() {
6868

6969
void setLatestConsumerGroupOffset(long consumerOffset);
7070

71+
void checkChannelStatusAndLogErrors(boolean tolerateErrors);
72+
7173
default CompletableFuture<Void> waitForLastProcessedRecordCommitted() {
7274
return CompletableFuture.completedFuture(null);
7375
}

‎src/main/java/com/snowflake/kafka/connector/internal/streaming/v2/SnowpipeStreamingPartitionChannel.java‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,10 @@
22

33
import static com.snowflake.kafka.connector.internal.SnowflakeErrors.ERROR_5027;
44
import static com.snowflake.kafka.connector.internal.SnowflakeErrors.ERROR_5028;
5+
import static com.snowflake.kafka.connector.internal.SnowflakeErrors.ERROR_5030;
56

67
import com.google.common.annotations.VisibleForTesting;
8+
import com.snowflake.ingest.streaming.ChannelStatus;
79
import com.snowflake.ingest.streaming.OpenChannelResult;
810
import com.snowflake.ingest.streaming.SFException;
911
import com.snowflake.ingest.streaming.SnowflakeStreamingIngestChannel;
@@ -583,6 +585,61 @@ public void setLatestConsumerGroupOffset(long consumerOffset) {
583585
}
584586
}
585587

588+
@Override
589+
public void checkChannelStatusAndLogErrors(final boolean tolerateErrors) {
590+
final ChannelStatus status = channel.getChannelStatus();
591+
logChannelStatus(status);
592+
handleChannelErrors(status, tolerateErrors);
593+
}
594+
595+
private void logChannelStatus(final ChannelStatus status) {
596+
LOGGER.info(
597+
"Channel status for channel=[{}]: databaseName=[{}], schemaName=[{}], pipeName=[{}],"
598+
+ " channelName=[{}], statusCode=[{}], latestCommittedOffsetToken=[{}],"
599+
+ " createdOn=[{}], rowsInsertedCount=[{}], rowsParsedCount=[{}],"
600+
+ " rowsErrorCount=[{}], lastErrorOffsetTokenUpperBound=[{}],"
601+
+ " lastErrorMessage=[{}], lastErrorTimestamp=[{}],"
602+
+ " serverAvgProcessingLatency=[{}], lastRefreshedOn=[{}]",
603+
this.getChannelNameFormatV1(),
604+
status.getDatabaseName(),
605+
status.getSchemaName(),
606+
status.getPipeName(),
607+
status.getChannelName(),
608+
status.getStatusCode(),
609+
status.getLatestCommittedOffsetToken(),
610+
status.getCreatedOn(),
611+
status.getRowsInsertedCount(),
612+
status.getRowsParsedCount(),
613+
status.getRowsErrorCount(),
614+
status.getLastErrorOffsetTokenUpperBound(),
615+
status.getLastErrorMessage(),
616+
status.getLastErrorTimestamp(),
617+
status.getServerAvgProcessingLatency(),
618+
status.getLastRefreshedOn());
619+
}
620+
621+
private void handleChannelErrors(final ChannelStatus status, final boolean tolerateErrors) {
622+
final long rowsErrorCount = status.getRowsErrorCount();
623+
if (rowsErrorCount > 0) {
624+
final String errorMessage =
625+
String.format(
626+
"Channel [%s] has %d errors. Last error message: %s, last error timestamp: %s,"
627+
+ " last error offset token upper bound: %s",
628+
this.getChannelNameFormatV1(),
629+
rowsErrorCount,
630+
status.getLastErrorMessage(),
631+
status.getLastErrorTimestamp(),
632+
status.getLastErrorOffsetTokenUpperBound());
633+
634+
if (tolerateErrors) {
635+
LOGGER.warn(errorMessage);
636+
} else {
637+
this.telemetryServiceV2.reportKafkaConnectFatalError(errorMessage);
638+
throw ERROR_5030.getException(errorMessage);
639+
}
640+
}
641+
}
642+
586643
/**
587644
* Enum representing which Streaming API is invoking the fallback supplier. ({@link
588645
* #streamingApiFallbackSupplier(StreamingApiFallbackInvoker)})

‎src/test/java/com/snowflake/kafka/connector/internal/streaming/FakeSnowflakeStreamingIngestChannel.java‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ public class FakeSnowflakeStreamingIngestChannel
2222

2323
private volatile boolean closed;
2424
private String offsetToken;
25+
private ChannelStatus channelStatus;
2526

2627
public FakeSnowflakeStreamingIngestChannel(String pipeName, String channelName) {
2728
this.pipeName = pipeName;
@@ -100,7 +101,14 @@ public synchronized String getLatestCommittedOffsetToken() {
100101

101102
@Override
102103
public ChannelStatus getChannelStatus() {
103-
throw new UnsupportedOperationException();
104+
if (channelStatus == null) {
105+
throw new UnsupportedOperationException("ChannelStatus not configured for test");
106+
}
107+
return channelStatus;
108+
}
109+
110+
public void setChannelStatus(final ChannelStatus channelStatus) {
111+
this.channelStatus = channelStatus;
104112
}
105113

106114
@Override

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

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -856,4 +856,42 @@ public void testStreamingIngestionWithExactlyOnceSemanticsOverlappingOffsets() t
856856

857857
service2.closeAll();
858858
}
859+
860+
/**
861+
* Test that getOffset works correctly when there are no channel errors. This test verifies that
862+
* the channel status check in getOffset does not throw an exception when no errors are present.
863+
*/
864+
@Test
865+
public void testGetOffsetWithNoChannelErrors() throws Exception {
866+
SnowflakeSinkService service =
867+
StreamingSinkServiceBuilder.builder(conn, config)
868+
.withSinkTaskContext(new InMemorySinkTaskContext(Collections.singleton(topicPartition)))
869+
.build();
870+
service.startPartition(topicPartition);
871+
872+
Converter converter = buildJsonConverter();
873+
SchemaAndValue input =
874+
converter.toConnectData(topic, "{\"name\":\"test\"}".getBytes(StandardCharsets.UTF_8));
875+
876+
List<SinkRecord> records = new ArrayList<>();
877+
for (int i = 0; i < 5; i++) {
878+
records.add(
879+
new SinkRecord(
880+
topic, partition, Schema.STRING_SCHEMA, "key" + i, input.schema(), input.value(), i));
881+
}
882+
883+
service.insert(records);
884+
885+
// getOffset should work without throwing exceptions when there are no channel errors
886+
// The channel status check happens inside getOffset
887+
TestUtils.assertWithRetry(
888+
() -> {
889+
long offset = service.getOffset(topicPartition);
890+
return offset == 5;
891+
},
892+
5,
893+
20);
894+
895+
service.closeAll();
896+
}
859897
}

‎src/test/java/com/snowflake/kafka/connector/internal/streaming/v2/SnowpipeStreamingPartitionChannelStatusTest.java‎

Lines changed: 160 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,160 @@
1+
package com.snowflake.kafka.connector.internal.streaming.v2;
2+
3+
import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
4+
import static org.junit.jupiter.api.Assertions.assertThrows;
5+
import static org.junit.jupiter.api.Assertions.assertTrue;
6+
import static org.mockito.Mockito.mock;
7+
import static org.mockito.Mockito.when;
8+
9+
import com.snowflake.ingest.streaming.ChannelStatus;
10+
import com.snowflake.ingest.streaming.SnowflakeStreamingIngestChannel;
11+
import com.snowflake.kafka.connector.internal.SnowflakeKafkaConnectorException;
12+
import com.snowflake.kafka.connector.internal.streaming.FakeSnowflakeStreamingIngestChannel;
13+
import com.snowflake.kafka.connector.internal.telemetry.SnowflakeTelemetryService;
14+
import java.time.Duration;
15+
import java.time.Instant;
16+
import org.junit.jupiter.api.BeforeEach;
17+
import org.junit.jupiter.api.Test;
18+
19+
class SnowpipeStreamingPartitionChannelStatusTest {
20+
21+
private FakeSnowflakeStreamingIngestChannel fakeChannel;
22+
private SnowflakeTelemetryService telemetryService;
23+
private TestableSnowpipeStreamingPartitionChannel partitionChannel;
24+
25+
@BeforeEach
26+
void setUp() {
27+
fakeChannel = new FakeSnowflakeStreamingIngestChannel("testPipe", "testChannel");
28+
telemetryService = mock(SnowflakeTelemetryService.class);
29+
partitionChannel = new TestableSnowpipeStreamingPartitionChannel(fakeChannel, telemetryService);
30+
}
31+
32+
@Test
33+
void shouldThrowExceptionWhenErrorCountGreaterThanZeroAndTolerateErrorsFalse() {
34+
// Given
35+
ChannelStatus statusWithErrors = createChannelStatusWithErrors(5, "Test error message");
36+
fakeChannel.setChannelStatus(statusWithErrors);
37+
38+
// When/Then
39+
SnowflakeKafkaConnectorException exception =
40+
assertThrows(
41+
SnowflakeKafkaConnectorException.class,
42+
() -> partitionChannel.checkChannelStatusAndLogErrors(false));
43+
44+
assertTrue(exception.getMessage().contains("5 errors"));
45+
assertTrue(exception.getMessage().contains("Test error message"));
46+
}
47+
48+
@Test
49+
void shouldNotThrowExceptionWhenErrorCountGreaterThanZeroAndTolerateErrorsTrue() {
50+
// Given
51+
ChannelStatus statusWithErrors = createChannelStatusWithErrors(5, "Test error message");
52+
fakeChannel.setChannelStatus(statusWithErrors);
53+
54+
// When/Then
55+
assertDoesNotThrow(() -> partitionChannel.checkChannelStatusAndLogErrors(true));
56+
}
57+
58+
@Test
59+
void shouldNotThrowExceptionWhenNoErrors() {
60+
// Given
61+
ChannelStatus statusWithoutErrors = createChannelStatusWithErrors(0, null);
62+
fakeChannel.setChannelStatus(statusWithoutErrors);
63+
64+
// When/Then - should not throw regardless of tolerateErrors setting
65+
assertDoesNotThrow(() -> partitionChannel.checkChannelStatusAndLogErrors(false));
66+
assertDoesNotThrow(() -> partitionChannel.checkChannelStatusAndLogErrors(true));
67+
}
68+
69+
private ChannelStatus createChannelStatusWithErrors(
70+
final long rowsErrorCount, final String lastErrorMessage) {
71+
ChannelStatus status = mock(ChannelStatus.class);
72+
when(status.getDatabaseName()).thenReturn("testDatabase");
73+
when(status.getSchemaName()).thenReturn("testSchema");
74+
when(status.getPipeName()).thenReturn("testPipe");
75+
when(status.getChannelName()).thenReturn("testChannel");
76+
when(status.getStatusCode()).thenReturn("SUCCESS");
77+
when(status.getLatestCommittedOffsetToken()).thenReturn("100");
78+
when(status.getCreatedOn()).thenReturn(Instant.now());
79+
when(status.getRowsInsertedCount()).thenReturn(1000L);
80+
when(status.getRowsParsedCount()).thenReturn(1005L);
81+
when(status.getRowsErrorCount()).thenReturn(rowsErrorCount);
82+
when(status.getLastErrorOffsetTokenUpperBound()).thenReturn("95");
83+
when(status.getLastErrorMessage()).thenReturn(lastErrorMessage);
84+
when(status.getLastErrorTimestamp()).thenReturn(Instant.now());
85+
when(status.getServerAvgProcessingLatency()).thenReturn(Duration.ofMillis(50));
86+
when(status.getLastRefreshedOn()).thenReturn(Instant.now());
87+
return status;
88+
}
89+
90+
/**
91+
* A testable version of SnowpipeStreamingPartitionChannel that allows injecting the channel and
92+
* telemetry service directly.
93+
*/
94+
private static class TestableSnowpipeStreamingPartitionChannel {
95+
private final SnowflakeStreamingIngestChannel channel;
96+
private final SnowflakeTelemetryService telemetryServiceV2;
97+
98+
TestableSnowpipeStreamingPartitionChannel(
99+
final SnowflakeStreamingIngestChannel channel,
100+
final SnowflakeTelemetryService telemetryService) {
101+
this.channel = channel;
102+
this.telemetryServiceV2 = telemetryService;
103+
}
104+
105+
void checkChannelStatusAndLogErrors(final boolean tolerateErrors) {
106+
final ChannelStatus status = channel.getChannelStatus();
107+
logChannelStatus(status);
108+
handleChannelErrors(status, tolerateErrors);
109+
}
110+
111+
private void logChannelStatus(final ChannelStatus status) {
112+
// Just log - actual logging tested via INFO level in real implementation
113+
System.out.printf(
114+
"Channel status: databaseName=[%s], schemaName=[%s], pipeName=[%s],"
115+
+ " channelName=[%s], statusCode=[%s], latestCommittedOffsetToken=[%s],"
116+
+ " createdOn=[%s], rowsInsertedCount=[%d], rowsParsedCount=[%d],"
117+
+ " rowsErrorCount=[%d], lastErrorOffsetTokenUpperBound=[%s],"
118+
+ " lastErrorMessage=[%s], lastErrorTimestamp=[%s],"
119+
+ " serverAvgProcessingLatency=[%s], lastRefreshedOn=[%s]%n",
120+
status.getDatabaseName(),
121+
status.getSchemaName(),
122+
status.getPipeName(),
123+
status.getChannelName(),
124+
status.getStatusCode(),
125+
status.getLatestCommittedOffsetToken(),
126+
status.getCreatedOn(),
127+
status.getRowsInsertedCount(),
128+
status.getRowsParsedCount(),
129+
status.getRowsErrorCount(),
130+
status.getLastErrorOffsetTokenUpperBound(),
131+
status.getLastErrorMessage(),
132+
status.getLastErrorTimestamp(),
133+
status.getServerAvgProcessingLatency(),
134+
status.getLastRefreshedOn());
135+
}
136+
137+
private void handleChannelErrors(final ChannelStatus status, final boolean tolerateErrors) {
138+
final long rowsErrorCount = status.getRowsErrorCount();
139+
if (rowsErrorCount > 0) {
140+
final String errorMessage =
141+
String.format(
142+
"Channel [%s] has %d errors. Last error message: %s, last error timestamp: %s,"
143+
+ " last error offset token upper bound: %s",
144+
channel.getFullyQualifiedChannelName(),
145+
rowsErrorCount,
146+
status.getLastErrorMessage(),
147+
status.getLastErrorTimestamp(),
148+
status.getLastErrorOffsetTokenUpperBound());
149+
150+
if (tolerateErrors) {
151+
System.out.println("WARN: " + errorMessage);
152+
} else {
153+
telemetryServiceV2.reportKafkaConnectFatalError(errorMessage);
154+
throw new SnowflakeKafkaConnectorException(errorMessage, "5030");
155+
}
156+
}
157+
}
158+
}
159+
}
160+

0 commit comments

Comments
 (0)