fix: update error when useGeneratedKeys

src/main/java/com/jd/jdbc/vitess/VitessStatement.java

@@ -99,6 +99,8 @@ public class VitessStatement extends AbstractVitessStatement {
 
     private static final String FUNCATION_IDENTITY = "@@IDENTITY";
 
+    private static Query.Field[] generatedKeyField;
+
     @Getter
     protected final Executor executor;
 
@@ -582,18 +584,15 @@ protected VtResultSet getGeneratedKeysInternal(long numKeys) throws SQLException
             throw new SQLException("does not support insert multiple rows in one sql statement");
         }
         long generatedKey = lastInsertId;
-        VtResultSet vtStaticResultSet = null;
-        Query.Field field = Query.Field.newBuilder().setName("GENERATED_KEY").setJdbcClassName("java.math.BigInteger").setType(Query.Type.UINT64).setColumnLength(20).setPrecision(20).build();
         List<List<VtResultValue>> rows = new ArrayList<>();
-
+        VtResultSet vtStaticResultSet = new VtResultSet(getGeneratedKeyField(), rows);
         if (!this.resultSets.isEmpty()) {
             if (generatedKey < 0) {
                 throw new SQLException("generatedKey error");
             }
             if (generatedKey != 0 && (numKeys > 0)) {
                 List<VtResultValue> row = Collections.singletonList(VtResultValue.newVtResultValue(Query.Type.UINT64, BigInteger.valueOf(generatedKey)));
                 rows.add(row);
-                vtStaticResultSet = new VtResultSet(new Query.Field[] {field}, rows);
             }
         }
         return vtStaticResultSet;
@@ -1145,6 +1144,20 @@ private void errorCount(String sql, Map<String, BindVariable> bindVariableMap, S
         StatementCollector.getStatementErrorCounter().labels(connection.getDefaultKeyspace(), VitessJdbcProperyUtil.getRole(connection.getProperties())).inc();
     }
 
+    private Query.Field[] getGeneratedKeyField() {
+        if (generatedKeyField == null) {
+            synchronized (VitessStatement.class) {
+                if (generatedKeyField == null) {
+                    Query.Field field = Query.Field.newBuilder().setName("GENERATED_KEY").setJdbcClassName("java.math.BigInteger")
+                        .setType(Query.Type.UINT64).setColumnLength(20).setPrecision(20).build();
+                    generatedKeyField = new Query.Field[] {field};
+                    return generatedKeyField;
+                }
+            }
+        }
+        return generatedKeyField;
+    }
+
     @Getter
     @AllArgsConstructor
     public static class ParseResult {

src/test/java/com/jd/jdbc/table/TableAutoGeneratedKeysTest.java

@@ -25,27 +25,27 @@
 import org.junit.After;
 import org.junit.Assert;
 import org.junit.Before;
-import org.junit.Rule;
 import org.junit.Test;
-import org.junit.rules.ExpectedException;
 import testsuite.internal.TestSuiteShardSpec;
 
 public class TableAutoGeneratedKeysTest extends VitessAutoGeneratedKeysTest {
 
-    @Rule
-    public ExpectedException thrown = ExpectedException.none();
-
     @Before
-    public void init() throws Exception {
+    @Override
+    public void testLoadDriver() throws Exception {
         getConn();
         TableTestUtil.setSplitTableConfig("table/autoGeneratedKeys.yml");
+        clean();
 
         sql1 = "insert into table_auto (id,ai,email) values(1,1,'x')";
         sql100 = "insert into table_auto (id,ai,email) values(1,100,'x')";
         sqlx = "insert into table_auto (id,ai,email) values(?,?,'x')";
         sqld = "insert into table_auto (id,ai,email) values(%d,%d,'x')";
         updateSql = "update table_auto set email = 'zz' where id = %d";
         deleteSql = "delete from table_auto where id = %d";
+        updateSql100 = "update table_auto set email = 'zz' where id = 1";
+        sql200 = "insert into table_auto (id,ai,email) values(200,200,'x')";
+        updatesql200 = "update table_auto set email = 'zz' where id = ?";
     }
 
     protected void getConn() throws SQLException {
@@ -55,9 +55,7 @@ protected void getConn() throws SQLException {
 
     @After
     public void close() throws Exception {
-        if (conn != null) {
-            conn.close();
-        }
+        closeConnection(conn);
         TableTestUtil.setDefaultTableConfig();
     }
 

src/test/java/com/jd/jdbc/table/TableRewriteBatchedStatementsTest.java

@@ -45,6 +45,14 @@ public void test12StatementExecuteBatch() throws Exception {
         super.test12StatementExecuteBatch();
     }
 
+    @Test
+    @Override
+    public void test12StatementUpdateExecuteBatch() throws Exception {
+        thrown.expect(java.sql.SQLFeatureNotSupportedException.class);
+        thrown.expectMessage("unsupported multiquery");
+        super.test12StatementUpdateExecuteBatch();
+    }
+
     @Test
     @Override
     public void test20PreparedStatementExecuteBatch() throws Exception {
@@ -70,8 +78,9 @@ public void test22PreparedStatementExecuteBatch() throws Exception {
     }
 
     @Test
-    @Override
-    public void test23SetNull() throws Exception {
-        super.test23SetNull();
+    public void test22PreparedStatementUpdateExecuteBatch() throws Exception {
+        thrown.expect(java.sql.SQLFeatureNotSupportedException.class);
+        thrown.expectMessage("unsupported multiquery");
+        super.test22PreparedStatementUpdateExecuteBatch();
     }
 }