[pulsar-io] feat: implement pip-297 for jdbc sinks

Author: 0asys

pulsar-io/jdbc/core/src/main/java/org/apache/pulsar/io/jdbc/JdbcAbstractSink.java

@@ -35,6 +35,7 @@
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import lombok.AllArgsConstructor;
@@ -50,6 +51,11 @@
  */
 @Slf4j
 public abstract class JdbcAbstractSink<T> implements Sink<T> {
+
+    private enum State {
+        OPEN, FAILED, CLOSED
+    }
+
     // ----- Runtime fields
     protected JdbcSinkConfig jdbcSinkConfig;
     @Getter
@@ -73,9 +79,12 @@ public abstract class JdbcAbstractSink<T> implements Sink<T> {
     private AtomicBoolean isFlushing;
     private int batchSize;
     private ScheduledExecutorService flushExecutor;
+    private SinkContext sinkContext;
+    private final AtomicReference<State> state = new AtomicReference<>(State.OPEN);
 
     @Override
     public void open(Map<String, Object> config, SinkContext sinkContext) throws Exception {
+        this.sinkContext = sinkContext;
         jdbcSinkConfig = JdbcSinkConfig.load(config, sinkContext);
         jdbcSinkConfig.validate();
 
@@ -148,6 +157,7 @@ private static List<String> getListFromConfig(String jdbcSinkConfig) {
 
     @Override
     public void close() throws Exception {
+        state.set(State.CLOSED);
         if (flushExecutor != null) {
             int timeoutMs = jdbcSinkConfig.getTimeoutMs() * 2;
             flushExecutor.shutdown();
@@ -310,8 +320,9 @@ private void flush() {
                         connection.rollback();
                     }
                 } catch (Exception ex) {
-                    throw new RuntimeException(ex);
+                    log.error("Failed to rollback transaction", ex);
                 }
+                fatal(e);
             }
 
             isFlushing.set(false);
@@ -385,4 +396,16 @@ private static boolean isBatchItemFailed(int returnCode) {
         return true;
     }
 
+    /**
+     * Signal a fatal exception to the framework.
+     * This will cause the function instance to terminate properly.
+     *
+     * @param e the fatal exception
+     */
+    private void fatal(Exception e) {
+        if (sinkContext != null && state.compareAndSet(State.OPEN, State.FAILED)) {
+            sinkContext.fatal(e);
+        }
+    }
+
 }

pulsar-io/jdbc/sqlite/src/test/java/org/apache/pulsar/io/jdbc/SqliteJdbcSinkTest.java

@@ -18,8 +18,10 @@
  */
 package org.apache.pulsar.io.jdbc;
 
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Maps;
@@ -56,6 +58,7 @@
 import org.apache.pulsar.common.schema.SchemaType;
 import org.apache.pulsar.functions.api.Record;
 import org.apache.pulsar.functions.source.PulsarRecord;
+import org.apache.pulsar.io.core.SinkContext;
 import org.awaitility.Awaitility;
 import org.testng.Assert;
 import org.testng.annotations.AfterMethod;
@@ -860,6 +863,50 @@ public void testNullValueAction(NullValueActionTestConfig config) throws Excepti
         }
     }
 
+    /**
+     * Test that fatal() is called when an unrecoverable exception occurs during flush.
+     * This verifies the PIP-297 implementation for proper termination of the sink.
+     */
+    @Test
+    public void testFatalCalledOnFlushException() throws Exception {
+        jdbcSink.close();
+        jdbcSink = null;
+
+        String jdbcUrl = sqliteUtils.sqliteUri();
+        Map<String, Object> conf = Maps.newHashMap();
+        conf.put("jdbcUrl", jdbcUrl);
+        conf.put("tableName", "nonexistent_table");  // This will cause an exception on flush
+        conf.put("key", "field3");
+        conf.put("nonKey", "field1,field2");
+        conf.put("batchSize", 1);
+
+        SinkContext mockSinkContext = mock(SinkContext.class);
+        AtomicReference<Throwable> fatalException = new AtomicReference<>();
+        doAnswer(invocation -> {
+            fatalException.set(invocation.getArgument(0));
+            return null;
+        }).when(mockSinkContext).fatal(any(Throwable.class));
+
+        SqliteJdbcAutoSchemaSink sinkWithContext = new SqliteJdbcAutoSchemaSink();
+        try {
+            sinkWithContext.open(conf, mockSinkContext);
+
+            Foo insertObj = new Foo("f1", "f2", 1);
+            Map<String, String> props = Maps.newHashMap();
+            props.put("ACTION", "INSERT");
+            CompletableFuture<Boolean> future = new CompletableFuture<>();
+            sinkWithContext.write(createMockFooRecord(insertObj, props, future));
+
+            // Wait for the flush to complete and fail
+            Awaitility.await().atMost(5, TimeUnit.SECONDS).untilAsserted(() -> {
+                verify(mockSinkContext).fatal(any(Throwable.class));
+                Assert.assertNotNull(fatalException.get());
+            });
+        } finally {
+            sinkWithContext.close();
+        }
+    }
+
     @SuppressWarnings("unchecked")
     private Record<GenericObject> createMockFooRecord(Foo record, Map<String, String> actionProperties,
                                                         CompletableFuture<Boolean> future) {