调整sr写入逻辑

dbswitch-product/dbswitch-product-starrocks/src/main/java/org/dromara/dbswitch/product/sr/FrontendEntity.java

@@ -1,13 +1,22 @@
+// Copyright tang.  All rights reserved.
+// https://gitee.com/inrgihc/dbswitch
+//
+// Use of this source code is governed by a BSD-style license
+//
+// Author: tang (inrgihc@126.com)
+// Date : 2020/1/2
+// Location: beijing , china
+/////////////////////////////////////////////////////////////
 package org.dromara.dbswitch.product.sr;
 
-
 import lombok.Data;
 
 @Data
 public class FrontendEntity {
-    String ip;
-    String httpport;
-    Boolean alive;
-    Boolean join;
-    String role;
+
+  private String ip;
+  private String httpport;
+  private Boolean alive;
+  private Boolean join;
+  private String role;
 }

dbswitch-product/dbswitch-product-starrocks/src/main/java/org/dromara/dbswitch/product/sr/StarRocksUtils.java

@@ -16,7 +16,18 @@ import cn.hutool.db.Entity;
 import cn.hutool.json.JSONArray;
 import cn.hutool.json.JSONObject;
 import cn.hutool.json.JSONUtil;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.sql.Timestamp;
+import java.util.ArrayList;
+import java.util.Base64;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+import javax.sql.DataSource;
 import lombok.extern.slf4j.Slf4j;
+import org.apache.commons.collections4.CollectionUtils;
 import org.apache.http.HttpHeaders;
 import org.apache.http.client.methods.CloseableHttpResponse;
 import org.apache.http.client.methods.HttpPut;
@@ -28,129 +39,118 @@ import org.apache.http.impl.client.HttpClients;
 import org.apache.http.util.EntityUtils;
 import org.dromara.dbswitch.common.entity.CloseableDataSource;
 
-import javax.sql.DataSource;
-import java.io.IOException;
-import java.nio.charset.StandardCharsets;
-import java.sql.Timestamp;
-import java.util.*;
-import java.util.stream.Collectors;
-
 @Slf4j
 public final class StarRocksUtils {
 
+  private String dbName;
+  private String tbName;
+  private String host;
+  private String username;
+  private String password;
+  private CloseableDataSource dataSource;
+  private String httpPort;
 
-    private String indexName;
-    private volatile String dbName;
-    private volatile String tbName;
-    private volatile String host;
-    private volatile String username;
-    private volatile String password;
-    private volatile CloseableDataSource dataSource;
-    private volatile String httpPort;
+  public void init(String schemaName, String tableName, DataSource dataSource) {
+    this.getHttpPort(dataSource);
+    this.dataSource = (CloseableDataSource) dataSource;
+    this.host = ReUtil.extractMulti("jdbc:mysql://(.*):[0-9]{2,8}/", this.dataSource.getJdbcUrl(), "$1");
+    this.username = this.dataSource.getUserName();
+    this.password = this.dataSource.getPassword();
+    this.tbName = tableName;
+    this.dbName = schemaName;
+  }
 
-
-    public void init(String schemaName, String tableName, DataSource dataSource) {
-        this.getHttpPort(dataSource);
-        this.dataSource = (CloseableDataSource) dataSource;
-        this.indexName = tableName;
-        this.host = ReUtil.extractMulti("jdbc:mysql://(.*):[0-9]{2,8}/", this.dataSource.getJdbcUrl(), "$1");
-        this.username = this.dataSource.getUserName();
-        this.password = this.dataSource.getPassword();
-        this.tbName = tableName;
-        this.dbName = schemaName;
+  public void getHttpPort(DataSource dataSource) {
+    Db use = Db.use(dataSource);
+    try {
+      List<Entity> frontends = use.query("SHOW FRONTENDS");
+      List<FrontendEntity> frontendEntities = BeanUtil.copyToList(frontends, FrontendEntity.class);
+      List<FrontendEntity> leader = frontendEntities.stream().filter(i -> i.getRole().equals("LEADER"))
+          .collect(Collectors.toList());
+      FrontendEntity frontendEntity = leader.get(0);
+      this.httpPort = frontendEntity.getHttpport();
+    } catch (Exception e) {
+      throw new RuntimeException(e);
     }
+  }
 
-    public void getHttpPort(DataSource dataSource) {
-        Db use = Db.use(dataSource);
-        try {
-            List<Entity> frontends = use.query("SHOW FRONTENDS");
-            List<FrontendEntity> frontendEntities = BeanUtil.copyToList(frontends, FrontendEntity.class);
-            List<FrontendEntity> leader = frontendEntities.stream().filter(i -> i.getRole().equals("LEADER")).collect(Collectors.toList());
-            FrontendEntity frontendEntity = leader.get(0);
-            this.httpPort = frontendEntity.getHttpport();
-        } catch (Exception e) {
-            log.error(e.getMessage());
+  public long addOrUpdateData(List<String> fieldNames, List<Object[]> recordValues) {
+    if (CollectionUtils.isEmpty(fieldNames) || CollectionUtils.isEmpty(recordValues)) {
+      return 0L;
+    }
+    List<Object> objectList = asObjectList(fieldNames, recordValues);
+    JSONArray array = JSONUtil.parseArray(objectList);
+    JSONObject jsonObject = JSONUtil.createObj().set("data", array);
+    try {
+      sendData(jsonObject.toString());
+      return recordValues.size();
+    } catch (Exception e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  private void sendData(String content) throws Exception {
+    final String loadUrl = String.format("http://%s:%s/api/%s/%s/_stream_load",
+        this.host,
+        this.httpPort,
+        this.dbName,
+        this.tbName);
+
+    final HttpClientBuilder httpClientBuilder = HttpClients
+        .custom()
+        .setRedirectStrategy(new DefaultRedirectStrategy() {
+          @Override
+          protected boolean isRedirectable(String method) {
+            return true;
+          }
+        });
+    try (CloseableHttpClient client = httpClientBuilder.build()) {
+      HttpPut put = new HttpPut(loadUrl);
+      StringEntity entity = new StringEntity(content, "UTF-8");
+      put.setHeader(HttpHeaders.EXPECT, "100-continue");
+      put.setHeader(HttpHeaders.AUTHORIZATION, basicAuthHeader(this.username, this.password));
+      put.setHeader("strip_outer_array", "true");
+      put.setHeader("format", "JSON");
+      put.setHeader("json_root", "$.data");
+      put.setHeader("ignore_json_size", "true");
+      put.setHeader("Content-Type", "application/json");
+      put.setEntity(entity);
+      try (CloseableHttpResponse response = client.execute(put)) {
+        String loadResult = "";
+        if (response.getEntity() != null) {
+          loadResult = EntityUtils.toString(response.getEntity());
         }
-    }
-
-    public long addOrUpdateData(List<String> fieldNames, List<Object[]> recordValues) {
-        List<Object> objectList = asObjectList(fieldNames, recordValues);
-        JSONArray array = JSONUtil.parseArray(objectList);
-        JSONObject jsonObject = JSONUtil.createObj()
-                .set("data", array);
-        try {
-            sendData(jsonObject.toString());
-            return recordValues.size();
-        } catch (Exception e) {
-            throw new RuntimeException(e);
+        final int statusCode = response.getStatusLine().getStatusCode();
+        // statusCode 200 just indicates that starrocks be service is ok, not stream load
+        // you should see the output content to find whether stream load is success
+        if (statusCode != 200) {
+          throw new IOException(
+              String.format("Stream load failed, statusCode=%s load result=%s", statusCode, loadResult));
         }
+      }
     }
+  }
 
-    private void sendData(String content) throws Exception {
+  private String basicAuthHeader(String username, String password) {
+    final String tobeEncode = username + ":" + password;
+    byte[] encoded = Base64.getEncoder().encode(tobeEncode.getBytes(StandardCharsets.UTF_8));
+    return "Basic " + new String(encoded);
+  }
 
-        final String loadUrl = String.format("http://%s:%s/api/%s/%s/_stream_load",
-                this.host,
-                this.httpPort,
-                this.dbName,
-                this.tbName);
-
-        final HttpClientBuilder httpClientBuilder = HttpClients
-                .custom()
-                .setRedirectStrategy(new DefaultRedirectStrategy() {
-                    @Override
-                    protected boolean isRedirectable(String method) {
-                        return true;
-                    }
-                });
-
-        try (CloseableHttpClient client = httpClientBuilder.build()) {
-            HttpPut put = new HttpPut(loadUrl);
-            StringEntity entity = new StringEntity(content, "UTF-8");
-            put.setHeader(HttpHeaders.EXPECT, "100-continue");
-            put.setHeader(HttpHeaders.AUTHORIZATION, basicAuthHeader(this.username, this.password));
-            put.setHeader("strip_outer_array", "true");
-            put.setHeader("format", "JSON");
-            put.setHeader("json_root", "$.data");
-            put.setHeader("ignore_json_size", "true");
-            put.setHeader("Content-Type", "application/json");
-            put.setEntity(entity);
-            try (CloseableHttpResponse response = client.execute(put)) {
-                String loadResult = "";
-                if (response.getEntity() != null) {
-                    loadResult = EntityUtils.toString(response.getEntity());
-                }
-                final int statusCode = response.getStatusLine().getStatusCode();
-                // statusCode 200 just indicates that starrocks be service is ok, not stream load
-                // you should see the output content to find whether stream load is success
-                if (statusCode != 200) {
-                    throw new IOException(
-                            String.format("Stream load failed, statusCode=%s load result=%s", statusCode, loadResult));
-                }
-            }
+  private List<Object> asObjectList(List<String> fieldNames, List<Object[]> recordValues) {
+    int fieldCount = Math.min(fieldNames.size(), recordValues.get(0).length);
+    List<Object> rows = new ArrayList<>(recordValues.size());
+    for (Object[] row : recordValues) {
+      Map<String, Object> columns = new LinkedHashMap<>(fieldCount);
+      for (int i = 0; i < fieldCount; ++i) {
+        Object rowValue = row[i];
+        if (row[i] instanceof Timestamp) {
+          rowValue = String.valueOf(rowValue);
         }
+        columns.put(fieldNames.get(i), rowValue);
+      }
+      rows.add(columns);
     }
-
-    private String basicAuthHeader(String username, String password) {
-        final String tobeEncode = username + ":" + password;
-        byte[] encoded = Base64.getEncoder().encode(tobeEncode.getBytes(StandardCharsets.UTF_8));
-        return "Basic " + new String(encoded);
-    }
-
-
-    private List<Object> asObjectList(List<String> fieldNames, List<Object[]> recordValues) {
-        int fieldCount = Math.min(fieldNames.size(), recordValues.get(0).length);
-        List<Object> rows = new ArrayList<>(recordValues.size());
-        for (Object[] row : recordValues) {
-            Map<String, Object> columns = new LinkedHashMap<>(fieldCount);
-            for (int i = 0; i < fieldCount; ++i) {
-                Object rowValue = row[i];
-                if (row[i] instanceof Timestamp) {
-                    rowValue = String.valueOf(rowValue);
-                }
-                columns.put(fieldNames.get(i), rowValue);
-            }
-            rows.add(columns);
-        }
-        return rows;
-    }
+    return rows;
+  }
 }

dbswitch-product/dbswitch-product-starrocks/src/main/java/org/dromara/dbswitch/product/sr/StarrocksFactoryProvider.java

@@ -9,16 +9,14 @@
 /////////////////////////////////////////////////////////////
 package org.dromara.dbswitch.product.sr;
 
-import org.dromara.dbswitch.core.annotation.Product;
+import javax.sql.DataSource;
 import org.dromara.dbswitch.common.type.ProductTypeEnum;
+import org.dromara.dbswitch.core.annotation.Product;
 import org.dromara.dbswitch.core.features.ProductFeatures;
 import org.dromara.dbswitch.core.provider.AbstractFactoryProvider;
 import org.dromara.dbswitch.core.provider.meta.MetadataProvider;
-import org.dromara.dbswitch.core.provider.sync.AutoCastTableDataSynchronizeProvider;
 import org.dromara.dbswitch.core.provider.sync.TableDataSynchronizeProvider;
-import org.dromara.dbswitch.core.provider.write.AutoCastTableDataWriteProvider;
 import org.dromara.dbswitch.core.provider.write.TableDataWriteProvider;
-import javax.sql.DataSource;
 
 @Product(ProductTypeEnum.STARROCKS)
 public class StarrocksFactoryProvider extends AbstractFactoryProvider {
@@ -47,5 +45,4 @@ public class StarrocksFactoryProvider extends AbstractFactoryProvider {
     return new StarrocksTableDataSynchronizer(this);
   }
 
-
 }

dbswitch-product/dbswitch-product-starrocks/src/main/java/org/dromara/dbswitch/product/sr/StarrocksMetadataQueryProvider.java

@@ -48,6 +48,11 @@ public class StarrocksMetadataQueryProvider extends AbstractMetadataProvider {
   private static final String QUERY_TABLE_METADATA_SQL =
       "SELECT `TABLE_COMMENT`,`TABLE_TYPE` FROM `information_schema`.`TABLES` "
           + "WHERE `TABLE_SCHEMA` = ? AND `TABLE_NAME` = ?";
+  private static final String QUERY_TABLE_PRIMARY_KEY_SQL = " SELECT COLUMN_NAME \n" +
+      " from information_schema.columns\n" +
+      " where TABLE_SCHEMA=? and TABLE_NAME=?\n" +
+      " and TABLE_CATALOG is null\n" +
+      " and COLUMN_KEY = 'PRI'";
 
   public StarrocksMetadataQueryProvider(ProductFactoryProvider factoryProvider) {
     super(factoryProvider);
@@ -166,21 +171,16 @@ public class StarrocksMetadataQueryProvider extends AbstractMetadataProvider {
 
   @Override
   public List<String> queryTablePrimaryKeys(Connection connection, String schemaName, String tableName) {
-    List<String> ret = new ArrayList<>();
-    try {
-      Statement statement = connection.createStatement();
-      String sql = String.format(" SELECT * \n" +
-                      " from information_schema.columns\n" +
-                      " where TABLE_SCHEMA=\"%s\" and TABLE_NAME=\"%s\"\n" +
-                      " and TABLE_CATALOG is null\n" +
-                      " and COLUMN_KEY = 'PRI';",
-              schemaName, tableName
-      );
-      ResultSet primaryKeys = statement.executeQuery(sql);
-      while (primaryKeys.next()) {
-        ret.add(primaryKeys.getString("COLUMN_NAME"));
+    try (PreparedStatement statement = connection.prepareStatement(QUERY_TABLE_PRIMARY_KEY_SQL)) {
+      statement.setString(1, schemaName);
+      statement.setString(2, tableName);
+      try (ResultSet primaryKeys = statement.executeQuery()) {
+        List<String> ret = new ArrayList<>();
+        while (primaryKeys.next()) {
+          ret.add(primaryKeys.getString(1));
+        }
+        return ret.stream().distinct().collect(Collectors.toList());
       }
-      return ret.stream().distinct().collect(Collectors.toList());
     } catch (SQLException e) {
       throw new RuntimeException(e);
     }
@@ -323,7 +323,10 @@ public class StarrocksMetadataQueryProvider extends AbstractMetadataProvider {
         break;
       case ColumnMetaData.TYPE_STRING:
         //see: https://docs.starrocks.io/zh/docs/category/string/
-        if (length <= 65533) {
+        long newLength = length * 3;
+        if (newLength < 255) {
+          retval += "VARCHAR(" + newLength + ")";
+        } else if (newLength <= 65533) {
           retval += "STRING";
         } else if (newLength <= 1048576) {
           retval += "VARCHAR(" + newLength + ")";

dbswitch-product/dbswitch-product-starrocks/src/main/java/org/dromara/dbswitch/product/sr/StarrocksTableDataSynchronizer.java

@@ -9,100 +9,54 @@
 /////////////////////////////////////////////////////////////
 package org.dromara.dbswitch.product.sr;
 
-
-import org.apache.commons.collections4.CollectionUtils;
+import java.util.List;
+import lombok.extern.slf4j.Slf4j;
 import org.dromara.dbswitch.common.entity.CloseableDataSource;
 import org.dromara.dbswitch.core.provider.ProductFactoryProvider;
-import org.dromara.dbswitch.core.provider.sync.DefaultTableDataSynchronizeProvider;
+import org.dromara.dbswitch.core.provider.sync.AutoCastTableDataSynchronizeProvider;
 
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Map;
+@Slf4j
+public class StarrocksTableDataSynchronizer extends AutoCastTableDataSynchronizeProvider {
 
-public class StarrocksTableDataSynchronizer extends DefaultTableDataSynchronizeProvider {
+  private List<String> fieldNames;
+  private final CloseableDataSource dataSource;
 
-    private volatile List<String> fieldNames;
-    private final CloseableDataSource dataSource;
+  private final StarRocksUtils starRocksUtils = new StarRocksUtils();
 
-    private final StarRocksUtils starRocksUtils = new StarRocksUtils();
+  public StarrocksTableDataSynchronizer(ProductFactoryProvider factoryProvider) {
+    super(factoryProvider);
+    dataSource = (CloseableDataSource) factoryProvider.getDataSource();
+  }
 
-
-    public StarrocksTableDataSynchronizer(ProductFactoryProvider factoryProvider) {
-        super(factoryProvider);
-        dataSource = (CloseableDataSource) factoryProvider.getDataSource();
+  @Override
+  public void prepare(String schemaName, String tableName, List<String> fieldNames, List<String> pks) {
+    this.fieldNames = fieldNames;
+    super.prepare(schemaName, tableName, fieldNames, pks);
+    try {
+      starRocksUtils.init(schemaName, tableName, dataSource);
+    } catch (Exception e) {
+      log.warn("Failed to init by StarRocksUtils#init(),information:: {}", e.getMessage());
     }
+  }
 
-    @Override
-    public void prepare(String schemaName, String tableName, List<String> fieldNames, List<String> pks) {
-        starRocksUtils.init(schemaName, tableName, dataSource);
-
-        this.fieldNames = fieldNames;
-
-
-        if (fieldNames.isEmpty() || pks.isEmpty() || fieldNames.size() < pks.size()) {
-            throw new IllegalArgumentException("字段列表和主键列表不能为空，或者字段总个数应不小于主键总个数");
-        }
-        if (!fieldNames.containsAll(pks)) {
-            throw new IllegalArgumentException("字段列表必须包含主键列表");
-        }
-
-        Map<String, Integer> columnType = getTableColumnMetaData(schemaName, tableName, fieldNames);
-        this.fieldOrders = new ArrayList<>(fieldNames);
-        this.pksOrders = new ArrayList<>(pks);
-
-        this.insertStatementSql = getInsertPrepareStatementSql(schemaName, tableName, fieldNames);
-        this.updateStatementSql = getUpdatePrepareStatementSql(schemaName, tableName, fieldNames, pks);
-        this.deleteStatementSql = getDeletePrepareStatementSql(schemaName, tableName, pks);
-
-        insertArgsType = new int[fieldNames.size()];
-        for (int k = 0; k < fieldNames.size(); ++k) {
-            String field = fieldNames.get(k);
-            insertArgsType[k] = columnType.get(field);
-        }
-
-        updateArgsType = new int[fieldNames.size()];
-        int idx = 0;
-        for (int i = 0; i < fieldNames.size(); ++i) {
-            String field = fieldNames.get(i);
-            if (!pks.contains(field)) {
-                updateArgsType[idx++] = columnType.get(field);
-            }
-        }
-        for (String pk : pks) {
-            updateArgsType[idx++] = columnType.get(pk);
-        }
-
-        deleteArgsType = new int[pks.size()];
-        for (int j = 0; j < pks.size(); ++j) {
-            String pk = pks.get(j);
-            deleteArgsType[j] = columnType.get(pk);
-        }
+  @Override
+  public long executeInsert(List<Object[]> recordValues) {
+    try {
+      return starRocksUtils.addOrUpdateData(fieldNames, recordValues);
+    } catch (Exception e) {
+      log.warn("Failed to addOrUpdateData by StarRocksUtils#addOrUpdateData(),information:: {}", e.getMessage());
+      return super.executeInsert(recordValues);
     }
+  }
 
-
-    @Override
-    public long executeInsert(List<Object[]> recordValues) {
-        if (CollectionUtils.isEmpty(fieldNames) || CollectionUtils.isEmpty(recordValues)) {
-            return 0L;
-        }
-        if (CollectionUtils.isEmpty(fieldNames) || CollectionUtils.isEmpty(recordValues)) {
-            return 0L;
-        }
-
-        return starRocksUtils.addOrUpdateData(fieldNames, recordValues);
+  @Override
+  public long executeUpdate(List<Object[]> recordValues) {
+    try {
+      return starRocksUtils.addOrUpdateData(fieldNames, recordValues);
+    } catch (Exception e) {
+      log.warn("Failed to addOrUpdateData by StarRocksUtils#addOrUpdateData(),information:: {}", e.getMessage());
+      return super.executeUpdate(recordValues);
     }
-
-    @Override
-    public long executeUpdate(List<Object[]> recordValues) {
-        if (CollectionUtils.isEmpty(fieldNames) || CollectionUtils.isEmpty(recordValues)) {
-            return 0L;
-        }
-        if (CollectionUtils.isEmpty(fieldNames) || CollectionUtils.isEmpty(recordValues)) {
-            return 0L;
-        }
-
-        return starRocksUtils.addOrUpdateData(fieldNames, recordValues);
-    }
-
+  }
 
 }

dbswitch-product/dbswitch-product-starrocks/src/main/java/org/dromara/dbswitch/product/sr/StarrocksTableDataWriteProvider.java

@@ -9,42 +9,41 @@
 /////////////////////////////////////////////////////////////
 package org.dromara.dbswitch.product.sr;
 
+import java.util.List;
 import lombok.extern.slf4j.Slf4j;
-import org.apache.commons.collections4.CollectionUtils;
 import org.dromara.dbswitch.common.entity.CloseableDataSource;
 import org.dromara.dbswitch.core.provider.ProductFactoryProvider;
-import org.dromara.dbswitch.core.provider.write.DefaultTableDataWriteProvider;
-
-import java.util.List;
+import org.dromara.dbswitch.core.provider.write.AutoCastTableDataWriteProvider;
 
 @Slf4j
-public class StarrocksTableDataWriteProvider extends DefaultTableDataWriteProvider {
+public class StarrocksTableDataWriteProvider extends AutoCastTableDataWriteProvider {
 
+  private final CloseableDataSource dataSource;
+  private final StarRocksUtils starRocksUtils = new StarRocksUtils();
 
-    private final CloseableDataSource dataSource;
-    private final StarRocksUtils starRocksUtils = new StarRocksUtils();
-    ;
-
-    public StarrocksTableDataWriteProvider(ProductFactoryProvider factoryProvider) {
-        super(factoryProvider);
-        dataSource = (CloseableDataSource) factoryProvider.getDataSource();
+  public StarrocksTableDataWriteProvider(ProductFactoryProvider factoryProvider) {
+    super(factoryProvider);
+    dataSource = (CloseableDataSource) factoryProvider.getDataSource();
+  }
 
+  @Override
+  public void prepareWrite(String schemaName, String tableName, List<String> fieldNames) {
+    super.prepareWrite(schemaName, tableName, fieldNames);
+    try {
+      starRocksUtils.init(schemaName, tableName, dataSource);
+    } catch (Exception e) {
+      log.warn("Failed to init by StarRocksUtils#init(),information: {}", e.getMessage());
     }
+  }
 
-    @Override
-    public void prepareWrite(String schemaName, String tableName, List<String> fieldNames) {
-        starRocksUtils.init(schemaName, tableName, dataSource);
+  @Override
+  public long write(List<String> fieldNames, List<Object[]> recordValues) {
+    try {
+      return starRocksUtils.addOrUpdateData(fieldNames, recordValues);
+    } catch (Exception e) {
+      log.warn("Failed to insertOrUpdate data by StarRocksUtils#addOrUpdateData(),information: {}", e.getMessage());
+      return super.write(fieldNames, recordValues);
     }
-
-    @Override
-    public long write(List<String> fieldNames, List<Object[]> recordValues) {
-        if (CollectionUtils.isEmpty(fieldNames) || CollectionUtils.isEmpty(recordValues)) {
-            return 0L;
-        }
-
-        return starRocksUtils.addOrUpdateData(fieldNames, recordValues);
-
-    }
-    
+  }
 
 }