[Feature][Connector] add starrocks save_mode

seatunnel-connectors-v2/connector-starrocks/src/main/java/org/apache/seatunnel/connectors/seatunnel/starrocks/catalog/StarRocksCatalog.java

@@ -37,6 +37,7 @@
 import org.apache.seatunnel.common.exception.CommonErrorCodeDeprecated;
 import org.apache.seatunnel.common.utils.JdbcUrlUtil;
 import org.apache.seatunnel.connectors.seatunnel.starrocks.exception.StarRocksConnectorException;
+import org.apache.seatunnel.connectors.seatunnel.starrocks.sink.StarRocksSaveModeUtil;
 
 import org.apache.commons.lang3.StringUtils;
 
@@ -72,6 +73,7 @@ public class StarRocksCatalog implements Catalog {
     protected final String baseUrl;
     protected String defaultUrl;
     private final JdbcUrlUtil.UrlInfo urlInfo;
+    private final String template;
 
     private static final Set<String> SYS_DATABASES = new HashSet<>();
     private static final Logger LOG = LoggerFactory.getLogger(StarRocksCatalog.class);
@@ -81,10 +83,10 @@ public class StarRocksCatalog implements Catalog {
         SYS_DATABASES.add("_statistics_");
     }
 
-    public StarRocksCatalog(String catalogName, String username, String pwd, String defaultUrl) {
+    public StarRocksCatalog(
+            String catalogName, String username, String pwd, String defaultUrl, String template) {
 
         checkArgument(StringUtils.isNotBlank(username));
-        checkArgument(StringUtils.isNotBlank(pwd));
         checkArgument(StringUtils.isNotBlank(defaultUrl));
         urlInfo = JdbcUrlUtil.getUrlInfo(defaultUrl);
         this.baseUrl = urlInfo.getUrlWithoutDatabase();
@@ -95,6 +97,7 @@ public StarRocksCatalog(String catalogName, String username, String pwd, String
         this.catalogName = catalogName;
         this.username = username;
         this.pwd = pwd;
+        this.template = template;
     }
 
     @Override
@@ -208,13 +211,70 @@ public CatalogTable getTable(TablePath tablePath)
     @Override
     public void createTable(TablePath tablePath, CatalogTable table, boolean ignoreIfExists)
             throws TableAlreadyExistException, DatabaseNotExistException, CatalogException {
-        throw new UnsupportedOperationException();
+        this.createTable(
+                StarRocksSaveModeUtil.fillingCreateSql(
+                        template,
+                        table.getTableId().getDatabaseName(),
+                        table.getTableId().getTableName(),
+                        table.getTableSchema()));
     }
 
     @Override
     public void dropTable(TablePath tablePath, boolean ignoreIfNotExists)
             throws TableNotExistException, CatalogException {
-        throw new UnsupportedOperationException();
+        try (Connection conn = DriverManager.getConnection(defaultUrl, username, pwd)) {
+            if (ignoreIfNotExists) {
+                conn.createStatement().execute("DROP TABLE IF EXISTS " + tablePath.getFullName());
+            } else {
+                conn.createStatement()
+                        .execute(String.format("DROP TABLE %s", tablePath.getFullName()));
+            }
+        } catch (Exception e) {
+            throw new CatalogException(
+                    String.format("Failed listing database in catalog %s", catalogName), e);
+        }
+    }
+
+    public void truncateTable(TablePath tablePath, boolean ignoreIfNotExists)
+            throws TableNotExistException, CatalogException {
+        try (Connection conn = DriverManager.getConnection(defaultUrl, username, pwd)) {
+            if (ignoreIfNotExists) {
+                conn.createStatement()
+                        .execute(String.format("TRUNCATE TABLE  %s", tablePath.getFullName()));
+            }
+        } catch (Exception e) {
+            throw new CatalogException(
+                    String.format("Failed TRUNCATE TABLE in catalog %s", tablePath.getFullName()),
+                    e);
+        }
+    }
+
+    public void executeSql(TablePath tablePath, String sql) {
+        try (Connection connection = DriverManager.getConnection(defaultUrl, username, pwd)) {
+            try (PreparedStatement ps = connection.prepareStatement(sql)) {
+                // Will there exist concurrent drop for one table?
+                ps.execute();
+            } catch (SQLException e) {
+                throw new CatalogException(String.format("Failed executeSql error %s", sql), e);
+            }
+        } catch (Exception e) {
+            throw new CatalogException(String.format("Failed EXECUTE SQL in catalog %s", sql), e);
+        }
+    }
+
+    public boolean isExistsData(TablePath tablePath) {
+        try (Connection connection = DriverManager.getConnection(defaultUrl, username, pwd)) {
+            String sql = String.format("select * from %s limit 1", tablePath.getFullName());
+            PreparedStatement ps = connection.prepareStatement(sql);
+            ResultSet resultSet = ps.executeQuery();
+            if (resultSet == null) {
+                return false;
+            }
+            return resultSet.next();
+        } catch (SQLException e) {
+            throw new CatalogException(
+                    String.format("Failed Connection JDBC error %s", tablePath.getTableName()), e);
+        }
     }
 
     @Override

seatunnel-connectors-v2/connector-starrocks/src/main/java/org/apache/seatunnel/connectors/seatunnel/starrocks/catalog/StarRocksCatalogFactory.java

@@ -23,6 +23,7 @@
 import org.apache.seatunnel.api.table.factory.CatalogFactory;
 import org.apache.seatunnel.api.table.factory.Factory;
 import org.apache.seatunnel.connectors.seatunnel.starrocks.config.StarRocksOptions;
+import org.apache.seatunnel.connectors.seatunnel.starrocks.config.StarRocksSinkOptions;
 
 import com.google.auto.service.AutoService;
 
@@ -36,7 +37,8 @@ public Catalog createCatalog(String catalogName, ReadonlyConfig options) {
                 catalogName,
                 options.get(StarRocksOptions.USERNAME),
                 options.get(StarRocksOptions.PASSWORD),
-                options.get(StarRocksOptions.BASE_URL));
+                options.get(StarRocksOptions.BASE_URL),
+                options.get(StarRocksSinkOptions.SAVE_MODE_CREATE_TEMPLATE));
     }
 
     @Override

seatunnel-connectors-v2/connector-starrocks/src/main/java/org/apache/seatunnel/connectors/seatunnel/starrocks/config/SinkConfig.java

@@ -19,6 +19,7 @@
 
 import org.apache.seatunnel.api.configuration.ReadonlyConfig;
 import org.apache.seatunnel.api.sink.DataSaveMode;
+import org.apache.seatunnel.api.sink.SchemaSaveMode;
 
 import lombok.Getter;
 import lombok.Setter;
@@ -58,7 +59,9 @@ public enum StreamLoadFormat {
 
     private String saveModeCreateTemplate;
 
+    private SchemaSaveMode schemaSaveMode;
     private DataSaveMode dataSaveMode;
+    private String customSql;
 
     private int httpSocketTimeout;
 
@@ -91,7 +94,9 @@ public static SinkConfig of(ReadonlyConfig config) {
         config.getOptional(StarRocksSinkOptions.COLUMN_SEPARATOR)
                 .ifPresent(sinkConfig::setColumnSeparator);
         sinkConfig.setLoadFormat(config.get(StarRocksSinkOptions.LOAD_FORMAT));
-        sinkConfig.setDataSaveMode(config.get(StarRocksSinkOptions.SAVE_MODE));
+        sinkConfig.setSchemaSaveMode(config.get(StarRocksSinkOptions.SCHEMA_SAVE_MODE));
+        sinkConfig.setDataSaveMode(config.get(StarRocksSinkOptions.DATA_SAVE_MODE));
+        sinkConfig.setCustomSql(config.get(StarRocksSinkOptions.CUSTOM_SQL));
         sinkConfig.setHttpSocketTimeout(config.get(StarRocksSinkOptions.HTTP_SOCKET_TIMEOUT_MS));
         return sinkConfig;
     }

seatunnel-connectors-v2/connector-starrocks/src/main/java/org/apache/seatunnel/connectors/seatunnel/starrocks/config/StarRocksSinkOptions.java

@@ -19,12 +19,10 @@
 
 import org.apache.seatunnel.api.configuration.Option;
 import org.apache.seatunnel.api.configuration.Options;
-import org.apache.seatunnel.api.configuration.SingleChoiceOption;
 import org.apache.seatunnel.api.sink.DataSaveMode;
-import org.apache.seatunnel.api.sink.SupportSaveMode;
+import org.apache.seatunnel.api.sink.SchemaSaveMode;
 import org.apache.seatunnel.connectors.seatunnel.starrocks.config.SinkConfig.StreamLoadFormat;
 
-import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
 
@@ -131,10 +129,15 @@ public interface StarRocksSinkOptions {
                     .enumType(StreamLoadFormat.class)
                     .defaultValue(StreamLoadFormat.JSON)
                     .withDescription("");
-
-    SingleChoiceOption<DataSaveMode> SAVE_MODE =
-            Options.key(SupportSaveMode.DATA_SAVE_MODE_KEY)
-                    .singleChoice(DataSaveMode.class, Arrays.asList(DataSaveMode.APPEND_DATA))
+    Option<SchemaSaveMode> SCHEMA_SAVE_MODE =
+            Options.key("schema_save_mode")
+                    .enumType(SchemaSaveMode.class)
+                    .defaultValue(SchemaSaveMode.CREATE_SCHEMA_WHEN_NOT_EXIST)
+                    .withDescription("schema_save_mode");
+
+    Option<DataSaveMode> DATA_SAVE_MODE =
+            Options.key("data_save_mode")
+                    .enumType(DataSaveMode.class)
                     .defaultValue(DataSaveMode.APPEND_DATA)
                     .withDescription(
                             "Table structure and data processing methods that already exist on the target end");
@@ -144,4 +147,7 @@ public interface StarRocksSinkOptions {
                     .intType()
                     .defaultValue(3 * 60 * 1000)
                     .withDescription("Set http socket timeout, default is 3 minutes.");
+
+    Option<String> CUSTOM_SQL =
+            Options.key("custom_sql").stringType().noDefaultValue().withDescription("custom_sql");
 }

seatunnel-connectors-v2/connector-starrocks/src/main/java/org/apache/seatunnel/connectors/seatunnel/starrocks/sink/StarRocksSink.java

@@ -17,10 +17,14 @@
 
 package org.apache.seatunnel.connectors.seatunnel.starrocks.sink;
 
+import org.apache.seatunnel.api.configuration.ReadonlyConfig;
 import org.apache.seatunnel.api.sink.DataSaveMode;
+import org.apache.seatunnel.api.sink.DefaultSaveModeHandler;
 import org.apache.seatunnel.api.sink.SaveModeHandler;
+import org.apache.seatunnel.api.sink.SchemaSaveMode;
 import org.apache.seatunnel.api.sink.SinkWriter;
 import org.apache.seatunnel.api.sink.SupportSaveMode;
+import org.apache.seatunnel.api.table.catalog.Catalog;
 import org.apache.seatunnel.api.table.catalog.CatalogTable;
 import org.apache.seatunnel.api.table.catalog.TablePath;
 import org.apache.seatunnel.api.table.type.SeaTunnelRow;
@@ -38,62 +42,51 @@ public class StarRocksSink extends AbstractSimpleSink<SeaTunnelRow, Void>
 
     private SeaTunnelRowType seaTunnelRowType;
     private final SinkConfig sinkConfig;
-    private final DataSaveMode dataSaveMode;
-
+    private DataSaveMode dataSaveMode;
+    private SchemaSaveMode schemaSaveMode;
     private final CatalogTable catalogTable;
 
-    public StarRocksSink(SinkConfig sinkConfig, CatalogTable catalogTable) {
+    public StarRocksSink(
+            SinkConfig sinkConfig, CatalogTable catalogTable, ReadonlyConfig readonlyConfig) {
         this.sinkConfig = sinkConfig;
         this.seaTunnelRowType = catalogTable.getTableSchema().toPhysicalRowDataType();
         this.catalogTable = catalogTable;
         this.dataSaveMode = sinkConfig.getDataSaveMode();
+        this.schemaSaveMode = sinkConfig.getSchemaSaveMode();
     }
 
     @Override
     public String getPluginName() {
         return StarRocksCatalogFactory.IDENTIFIER;
     }
 
-    private void autoCreateTable(String template) {
-        StarRocksCatalog starRocksCatalog =
-                new StarRocksCatalog(
-                        "StarRocks",
-                        sinkConfig.getUsername(),
-                        sinkConfig.getPassword(),
-                        sinkConfig.getJdbcUrl());
-        if (!starRocksCatalog.databaseExists(sinkConfig.getDatabase())) {
-            starRocksCatalog.createDatabase(TablePath.of(sinkConfig.getDatabase(), ""), true);
-        }
-        if (!starRocksCatalog.tableExists(
-                TablePath.of(sinkConfig.getDatabase(), sinkConfig.getTable()))) {
-            starRocksCatalog.createTable(
-                    StarRocksSaveModeUtil.fillingCreateSql(
-                            template,
-                            sinkConfig.getDatabase(),
-                            sinkConfig.getTable(),
-                            catalogTable.getTableSchema()));
-        }
-    }
-
     @Override
     public AbstractSinkWriter<SeaTunnelRow, Void> createWriter(SinkWriter.Context context) {
         return new StarRocksSinkWriter(sinkConfig, seaTunnelRowType);
     }
 
     @Override
     public Optional<SaveModeHandler> getSaveModeHandler() {
-        return Optional.empty();
-    }
-
-    /*@Override
-    public DataSaveMode getUserConfigSaveMode() {
-        return dataSaveMode;
+        TablePath tablePath =
+                TablePath.of(
+                        catalogTable.getTableId().getDatabaseName(),
+                        catalogTable.getTableId().getSchemaName(),
+                        catalogTable.getTableId().getTableName());
+        Catalog catalog =
+                new StarRocksCatalog(
+                        "StarRocks",
+                        sinkConfig.getUsername(),
+                        sinkConfig.getPassword(),
+                        sinkConfig.getJdbcUrl(),
+                        sinkConfig.getSaveModeCreateTemplate());
+        catalog.open();
+        return Optional.of(
+                new DefaultSaveModeHandler(
+                        schemaSaveMode,
+                        dataSaveMode,
+                        catalog,
+                        tablePath,
+                        catalogTable,
+                        sinkConfig.getCustomSql()));
     }
-
-    @Override
-    public void handleSaveMode(DataSaveMode saveMode) {
-        if (catalogTable != null) {
-            autoCreateTable(sinkConfig.getSaveModeCreateTemplate());
-        }
-    }*/
 }

seatunnel-connectors-v2/connector-starrocks/src/main/java/org/apache/seatunnel/connectors/seatunnel/starrocks/sink/StarRocksSinkFactory.java

@@ -18,7 +18,9 @@
 package org.apache.seatunnel.connectors.seatunnel.starrocks.sink;
 
 import org.apache.seatunnel.api.configuration.util.OptionRule;
+import org.apache.seatunnel.api.sink.DataSaveMode;
 import org.apache.seatunnel.api.table.catalog.CatalogTable;
+import org.apache.seatunnel.api.table.catalog.TableIdentifier;
 import org.apache.seatunnel.api.table.connector.TableSink;
 import org.apache.seatunnel.api.table.factory.Factory;
 import org.apache.seatunnel.api.table.factory.TableSinkFactory;
@@ -31,6 +33,11 @@
 
 import com.google.auto.service.AutoService;
 
+import static org.apache.seatunnel.api.sink.SinkReplaceNameConstant.REPLACE_DATABASE_NAME_KEY;
+import static org.apache.seatunnel.api.sink.SinkReplaceNameConstant.REPLACE_SCHEMA_NAME_KEY;
+import static org.apache.seatunnel.api.sink.SinkReplaceNameConstant.REPLACE_TABLE_NAME_KEY;
+import static org.apache.seatunnel.connectors.seatunnel.starrocks.config.StarRocksSinkOptions.DATA_SAVE_MODE;
+
 @AutoService(Factory.class)
 public class StarRocksSinkFactory implements TableSinkFactory {
     @Override
@@ -54,9 +61,14 @@ public OptionRule optionRule() {
                         StarRocksSinkOptions.RETRY_BACKOFF_MULTIPLIER_MS,
                         StarRocksSinkOptions.STARROCKS_CONFIG,
                         StarRocksSinkOptions.ENABLE_UPSERT_DELETE,
-                        StarRocksSinkOptions.SAVE_MODE,
+                        StarRocksSinkOptions.SCHEMA_SAVE_MODE,
+                        StarRocksSinkOptions.DATA_SAVE_MODE,
                         StarRocksSinkOptions.SAVE_MODE_CREATE_TEMPLATE,
                         StarRocksSinkOptions.HTTP_SOCKET_TIMEOUT_MS)
+                .conditional(
+                        DATA_SAVE_MODE,
+                        DataSaveMode.CUSTOM_PROCESSING,
+                        StarRocksSinkOptions.CUSTOM_SQL)
                 .build();
     }
 
@@ -67,6 +79,47 @@ public TableSink createSink(TableSinkFactoryContext context) {
         if (StringUtils.isBlank(sinkConfig.getTable())) {
             sinkConfig.setTable(catalogTable.getTableId().getTableName());
         }
-        return () -> new StarRocksSink(sinkConfig, catalogTable);
+        // get source table relevant information
+        TableIdentifier tableId = catalogTable.getTableId();
+        String sourceDatabaseName = tableId.getDatabaseName();
+        String sourceSchemaName = tableId.getSchemaName();
+        String sourceTableName = tableId.getTableName();
+        // get sink table relevant information
+        String sinkDatabaseName = sinkConfig.getDatabase();
+        String sinkTableName = sinkConfig.getTable();
+        // to replace
+        String finalDatabaseName =
+                sinkDatabaseName.replace(REPLACE_DATABASE_NAME_KEY, sourceDatabaseName);
+        String finalTableName = this.replaceFullTableName(sinkTableName, tableId);
+        // rebuild TableIdentifier and catalogTable
+        TableIdentifier newTableId =
+                TableIdentifier.of(
+                        tableId.getCatalogName(), finalDatabaseName, null, finalTableName);
+        catalogTable =
+                CatalogTable.of(
+                        newTableId,
+                        catalogTable.getTableSchema(),
+                        catalogTable.getOptions(),
+                        catalogTable.getPartitionKeys(),
+                        catalogTable.getCatalogName());
+
+        CatalogTable finalCatalogTable = catalogTable;
+        // reset
+        sinkConfig.setTable(finalTableName);
+        sinkConfig.setDatabase(finalDatabaseName);
+        return () -> new StarRocksSink(sinkConfig, finalCatalogTable, context.getOptions());
+    }
+
+    private String replaceFullTableName(String original, TableIdentifier tableId) {
+        if (StringUtils.isNotBlank(tableId.getDatabaseName())) {
+            original = original.replace(REPLACE_DATABASE_NAME_KEY, tableId.getDatabaseName());
+        }
+        if (StringUtils.isNotBlank(tableId.getSchemaName())) {
+            original = original.replace(REPLACE_SCHEMA_NAME_KEY, tableId.getSchemaName());
+        }
+        if (StringUtils.isNotBlank(tableId.getTableName())) {
+            original = original.replace(REPLACE_TABLE_NAME_KEY, tableId.getTableName());
+        }
+        return original;
     }
 }

seatunnel-connectors-v2/connector-starrocks/src/test/java/org/apache/seatunnel/connectors/seatunnel/starrocks/catalog/StarRocksCatalogTest.java

@@ -36,7 +36,7 @@ public class StarRocksCatalogTest {
     public void testCatalog() {
         StarRocksCatalog catalog =
                 new StarRocksCatalog(
-                        "starrocks", "root", "123456", "jdbc:mysql://47.108.65.163:9030/");
+                        "starrocks", "root", "123456", "jdbc:mysql://47.108.65.163:9030/", "");
         List<String> databases = catalog.listDatabases();
         LOGGER.info("find databases: " + databases);