feat(sink): support jdbc sink async truncate

java/connector-node/connector-api/src/main/java/com/risingwave/connector/api/sink/SinkWriter.java

@@ -22,7 +22,7 @@
 public interface SinkWriter {
     void beginEpoch(long epoch);
 
-    void write(Iterable<SinkRow> rows);
+    boolean write(Iterable<SinkRow> rows);
 
     Optional<ConnectorServiceProto.SinkMetadata> barrier(boolean isCheckpoint);
 

java/connector-node/connector-api/src/main/java/com/risingwave/connector/api/sink/SinkWriterV1.java

@@ -43,11 +43,12 @@ public SinkWriterV1 getInner() {
         public void beginEpoch(long epoch) {}
 
         @Override
-        public void write(Iterable<SinkRow> rows) {
+        public boolean write(Iterable<SinkRow> rows) {
             if (!hasBegun) {
                 hasBegun = true;
             }
             this.inner.write(rows.iterator());
+            return false;
         }
 
         @Override

java/connector-node/risingwave-connector-service/src/main/java/com/risingwave/connector/SinkWriterStreamObserver.java

@@ -125,13 +125,29 @@ public void onNext(ConnectorServiceProto.SinkWriterStreamRequest sinkTask) {
                             .asRuntimeException();
                 }
 
+                boolean batchWritten;
+
                 try (CloseableIterable<SinkRow> rowIter = deserializer.deserialize(batch)) {
-                    sink.write(
-                            new MonitoredRowIterable(
-                                    rowIter, connectorName, String.valueOf(sinkId)));
+                    batchWritten =
+                            sink.write(
+                                    new MonitoredRowIterable(
+                                            rowIter, connectorName, String.valueOf(sinkId)));
                 }
 
                 currentBatchId = batch.getBatchId();
+
+                if (batchWritten) {
+                    responseObserver.onNext(
+                            ConnectorServiceProto.SinkWriterStreamResponse.newBuilder()
+                                    .setBatch(
+                                            ConnectorServiceProto.SinkWriterStreamResponse
+                                                    .BatchWrittenResponse.newBuilder()
+                                                    .setEpoch(currentEpoch)
+                                                    .setBatchId(currentBatchId)
+                                                    .build())
+                                    .build());
+                }
+
                 LOG.debug("Batch {} written to epoch {}", currentBatchId, batch.getEpoch());
             } else if (sinkTask.hasBarrier()) {
                 if (!isInitialized()) {

java/connector-node/risingwave-connector-test/src/test/java/com/risingwave/connector/sink/jdbc/JDBCSinkTest.java

@@ -16,7 +16,6 @@
 
 import static org.junit.Assert.*;
 
-import com.google.common.collect.Iterators;
 import com.google.common.collect.Lists;
 import com.risingwave.connector.JDBCSink;
 import com.risingwave.connector.JDBCSinkConfig;
@@ -26,6 +25,7 @@
 import com.risingwave.proto.Data.DataType.TypeName;
 import com.risingwave.proto.Data.Op;
 import java.sql.*;
+import java.util.Arrays;
 import org.junit.Test;
 import org.testcontainers.containers.JdbcDatabaseContainer;
 import org.testcontainers.containers.MySQLContainer;
@@ -84,7 +84,7 @@ static void testJDBCSync(JdbcDatabaseContainer<?> container, TestType testType)
         Connection conn = sink.getConn();
 
         sink.write(
-                Iterators.forArray(
+                Arrays.asList(
                         new ArraySinkRow(
                                 Op.INSERT,
                                 1,
@@ -94,7 +94,7 @@ static void testJDBCSync(JdbcDatabaseContainer<?> container, TestType testType)
                                 new Timestamp(1000000000),
                                 "{\"key\": \"password\", \"value\": \"Singularity123\"}",
                                 "I want to sleep".getBytes())));
-        sink.sync();
+        sink.barrier(true);
 
         Statement stmt = conn.createStatement();
         try (var rs = stmt.executeQuery(String.format("SELECT * FROM %s", tableName))) {
@@ -106,7 +106,7 @@ static void testJDBCSync(JdbcDatabaseContainer<?> container, TestType testType)
         }
 
         sink.write(
-                Iterators.forArray(
+                Arrays.asList(
                         new ArraySinkRow(
                                 Op.INSERT,
                                 2,
@@ -116,7 +116,7 @@ static void testJDBCSync(JdbcDatabaseContainer<?> container, TestType testType)
                                 new Timestamp(1000000000),
                                 "{\"key\": \"password\", \"value\": \"Singularity123\"}",
                                 "I want to sleep".getBytes())));
-        sink.sync();
+        sink.barrier(true);
         try (var rs = stmt.executeQuery(String.format("SELECT * FROM %s", tableName))) {
             int count;
             for (count = 0; rs.next(); ) {
@@ -126,7 +126,7 @@ static void testJDBCSync(JdbcDatabaseContainer<?> container, TestType testType)
         }
         stmt.close();
 
-        sink.sync();
+        sink.barrier(true);
         sink.drop();
     }
 
@@ -144,7 +144,7 @@ static void testJDBCWrite(JdbcDatabaseContainer<?> container, TestType testType)
         Statement stmt = conn.createStatement();
 
         sink.write(
-                Iterators.forArray(
+                Arrays.asList(
                         new ArraySinkRow(
                                 Op.INSERT,
                                 1,
@@ -171,7 +171,7 @@ static void testJDBCWrite(JdbcDatabaseContainer<?> container, TestType testType)
         }
 
         sink.write(
-                Iterators.forArray(
+                Arrays.asList(
                         new ArraySinkRow(
                                 Op.UPDATE_DELETE,
                                 1,
@@ -216,7 +216,7 @@ static void testJDBCWrite(JdbcDatabaseContainer<?> container, TestType testType)
             assertFalse(rs.next());
         }
 
-        sink.sync();
+        sink.barrier(true);
         stmt.close();
     }
 

java/connector-node/risingwave-sink-iceberg/src/main/java/com/risingwave/connector/AppendOnlyIcebergSinkWriter.java

@@ -52,7 +52,7 @@ public AppendOnlyIcebergSinkWriter(
     }
 
     @Override
-    public void write(Iterable<SinkRow> rows) {
+    public boolean write(Iterable<SinkRow> rows) {
         for (SinkRow row : rows) {
             switch (row.getOp()) {
                 case INSERT:
@@ -106,6 +106,7 @@ public void write(Iterable<SinkRow> rows) {
                             .asRuntimeException();
             }
         }
+        return false;
     }
 
     @Override

java/connector-node/risingwave-sink-iceberg/src/main/java/com/risingwave/connector/UpsertIcebergSinkWriter.java

@@ -139,7 +139,7 @@ private List<Comparable<Object>> getKeyFromRow(SinkRow row) {
     }
 
     @Override
-    public void write(Iterable<SinkRow> rows) {
+    public boolean write(Iterable<SinkRow> rows) {
         for (SinkRow row : rows) {
             if (row.size() != tableSchema.getColumnNames().length) {
                 throw Status.FAILED_PRECONDITION
@@ -188,6 +188,7 @@ public void write(Iterable<SinkRow> rows) {
                             .asRuntimeException();
             }
         }
+        return false;
     }
 
     @Override

java/connector-node/risingwave-sink-jdbc/src/main/java/com/risingwave/connector/JDBCSink.java

@@ -16,23 +16,26 @@
 
 import com.risingwave.connector.api.TableSchema;
 import com.risingwave.connector.api.sink.SinkRow;
-import com.risingwave.connector.api.sink.SinkWriterBase;
+import com.risingwave.connector.api.sink.SinkWriter;
 import com.risingwave.connector.jdbc.JdbcDialect;
+import com.risingwave.proto.ConnectorServiceProto;
 import com.risingwave.proto.Data;
 import io.grpc.Status;
 import java.sql.*;
 import java.util.*;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-public class JDBCSink extends SinkWriterBase {
+public class JDBCSink implements SinkWriter {
     private static final String ERROR_REPORT_TEMPLATE = "Error when exec %s, message %s";
 
     private final JdbcDialect jdbcDialect;
     private final JDBCSinkConfig config;
     private final Connection conn;
     private final List<String> pkColumnNames;
 
+    private final TableSchema tableSchema;
+
     public static final String JDBC_COLUMN_NAME_KEY = "COLUMN_NAME";
     public static final String JDBC_DATA_TYPE_KEY = "DATA_TYPE";
 
@@ -45,7 +48,7 @@ public class JDBCSink extends SinkWriterBase {
     private static final Logger LOG = LoggerFactory.getLogger(JDBCSink.class);
 
     public JDBCSink(JDBCSinkConfig config, TableSchema tableSchema) {
-        super(tableSchema);
+        this.tableSchema = tableSchema;
 
         var jdbcUrl = config.getJdbcUrl().toLowerCase();
         var factory = JdbcUtils.getDialectFactory(jdbcUrl);
@@ -175,7 +178,7 @@ private PreparedStatement prepareInsertStatement(SinkRow row) {
         }
         try {
             var preparedStmt = insertPreparedStmt;
-            jdbcDialect.bindInsertIntoStatement(preparedStmt, conn, getTableSchema(), row);
+            jdbcDialect.bindInsertIntoStatement(preparedStmt, conn, tableSchema, row);
             preparedStmt.addBatch();
             return preparedStmt;
         } catch (SQLException e) {
@@ -192,15 +195,15 @@ private PreparedStatement prepareUpsertStatement(SinkRow row) {
             var preparedStmt = upsertPreparedStmt;
             switch (row.getOp()) {
                 case INSERT:
-                    jdbcDialect.bindUpsertStatement(preparedStmt, conn, getTableSchema(), row);
+                    jdbcDialect.bindUpsertStatement(preparedStmt, conn, tableSchema, row);
                     break;
                 case UPDATE_INSERT:
                     if (!updateFlag) {
                         throw Status.FAILED_PRECONDITION
                                 .withDescription("an UPDATE_DELETE should precede an UPDATE_INSERT")
                                 .asRuntimeException();
                     }
-                    jdbcDialect.bindUpsertStatement(preparedStmt, conn, getTableSchema(), row);
+                    jdbcDialect.bindUpsertStatement(preparedStmt, conn, tableSchema, row);
                     updateFlag = false;
                     break;
                 default:
@@ -235,7 +238,7 @@ private PreparedStatement prepareDeleteStatement(SinkRow row) {
         try {
             int placeholderIdx = 1;
             for (String primaryKey : pkColumnNames) {
-                Object fromRow = getTableSchema().getFromRow(primaryKey, row);
+                Object fromRow = tableSchema.getFromRow(primaryKey, row);
                 deletePreparedStmt.setObject(placeholderIdx++, fromRow);
             }
             deletePreparedStmt.addBatch();
@@ -249,13 +252,15 @@ private PreparedStatement prepareDeleteStatement(SinkRow row) {
     }
 
     @Override
-    public void write(Iterator<SinkRow> rows) {
+    public void beginEpoch(long epoch) {}
+
+    @Override
+    public boolean write(Iterable<SinkRow> rows) {
         PreparedStatement deleteStatement = null;
         PreparedStatement upsertStatement = null;
         PreparedStatement insertStatement = null;
 
-        while (rows.hasNext()) {
-            SinkRow row = rows.next();
+        for (SinkRow row : rows) {
             if (row.getOp() == Data.Op.UPDATE_DELETE) {
                 updateFlag = true;
                 continue;
@@ -285,6 +290,7 @@ public void write(Iterator<SinkRow> rows) {
                             String.format(ERROR_REPORT_TEMPLATE, e.getSQLState(), e.getMessage()))
                     .asRuntimeException();
         }
+        return true;
     }
 
     private void executeStatement(PreparedStatement stmt) throws SQLException {
@@ -297,13 +303,14 @@ private void executeStatement(PreparedStatement stmt) throws SQLException {
     }
 
     @Override
-    public void sync() {
+    public Optional<ConnectorServiceProto.SinkMetadata> barrier(boolean isCheckpoint) {
         if (updateFlag) {
             throw Status.FAILED_PRECONDITION
                     .withDescription(
                             "expected UPDATE_INSERT to complete an UPDATE operation, got `sync`")
                     .asRuntimeException();
         }
+        return Optional.empty();
     }
 
     @Override

java/connector-node/risingwave-sink-jdbc/src/main/java/com/risingwave/connector/JDBCSinkFactory.java

@@ -19,7 +19,6 @@
 import com.risingwave.connector.api.TableSchema;
 import com.risingwave.connector.api.sink.SinkFactory;
 import com.risingwave.connector.api.sink.SinkWriter;
-import com.risingwave.connector.api.sink.SinkWriterV1;
 import com.risingwave.proto.Catalog.SinkType;
 import io.grpc.Status;
 import java.sql.*;
@@ -40,7 +39,7 @@ public class JDBCSinkFactory implements SinkFactory {
     public SinkWriter createWriter(TableSchema tableSchema, Map<String, String> tableProperties) {
         ObjectMapper mapper = new ObjectMapper();
         JDBCSinkConfig config = mapper.convertValue(tableProperties, JDBCSinkConfig.class);
-        return new SinkWriterV1.Adapter(new JDBCSink(config, tableSchema));
+        return new JDBCSink(config, tableSchema);
     }
 
     @Override

proto/connector_service.proto

@@ -88,9 +88,15 @@ message SinkWriterStreamResponse {
     SinkMetadata metadata = 2;
   }
 
+  message BatchWrittenResponse {
+    uint64 epoch = 1;
+    uint64 batch_id = 2;
+  }
+
   oneof response {
     StartResponse start = 1;
     CommitResponse commit = 2;
+    BatchWrittenResponse batch = 3;
   }
 }