bugfix: can not get generated keys value. (apache#2603)

dongYES · May 26, 2020 · 476416a · 476416a
1 parent 53963d4
commit 476416a
Show file tree

Hide file tree

Showing 6 changed files with 36 additions and 16 deletions.
diff --git a/rm-datasource/src/main/java/io/seata/rm/datasource/AbstractStatementProxy.java b/rm-datasource/src/main/java/io/seata/rm/datasource/AbstractStatementProxy.java
@@ -15,6 +15,8 @@
  */
 package io.seata.rm.datasource;
 
+import javax.sql.rowset.CachedRowSet;
+import javax.sql.rowset.RowSetProvider;
 import java.sql.Connection;
 import java.sql.ResultSet;
 import java.sql.SQLException;
@@ -40,6 +42,11 @@ public abstract class AbstractStatementProxy<T extends Statement> implements Sta
      */
     protected T targetStatement;
 
+    /**
+     * The generated keys cached row set.
+     */
+    private CachedRowSet generatedKeysRowSet;
+
     /**
      * The Target sql.
      */
@@ -242,7 +249,13 @@ public boolean getMoreResults(int current) throws SQLException {
 
     @Override
     public ResultSet getGeneratedKeys() throws SQLException {
-        return targetStatement.getGeneratedKeys();
+        if (generatedKeysRowSet != null) {
+            return generatedKeysRowSet;
+        }
+        ResultSet rs = targetStatement.getGeneratedKeys();
+        generatedKeysRowSet = RowSetProvider.newFactory().createCachedRowSet();
+        generatedKeysRowSet.populate(rs);
+        return generatedKeysRowSet;
     }
 
     @Override

diff --git a/rm-datasource/src/main/java/io/seata/rm/datasource/exec/BaseInsertExecutor.java b/rm-datasource/src/main/java/io/seata/rm/datasource/exec/BaseInsertExecutor.java
@@ -204,7 +204,7 @@ protected List<Object> parsePkValuesFromStatement() {
      */
     public List<Object> getGeneratedKeys() throws SQLException {
         // PK is just auto generated
-        ResultSet genKeys = statementProxy.getTargetStatement().getGeneratedKeys();
+        ResultSet genKeys = statementProxy.getGeneratedKeys();
         List<Object> pkValues = new ArrayList<>();
         while (genKeys.next()) {
             Object v = genKeys.getObject(1);
@@ -213,6 +213,11 @@ public List<Object> getGeneratedKeys() throws SQLException {
         if (pkValues.isEmpty()) {
             throw new NotSupportYetException(String.format("not support sql [%s]", sqlRecognizer.getOriginalSQL()));
         }
+        try {
+            genKeys.beforeFirst();
+        } catch (SQLException e) {
+            LOGGER.warn("Fail to reset ResultSet cursor. can not get primary key value");
+        }
         return pkValues;
     }
 

diff --git a/rm-datasource/src/main/java/io/seata/rm/datasource/exec/mysql/MySQLInsertExecutor.java b/rm-datasource/src/main/java/io/seata/rm/datasource/exec/mysql/MySQLInsertExecutor.java
@@ -81,7 +81,7 @@ public List<Object> getPkValuesByAuto() throws SQLException {
 
         ResultSet genKeys;
         try {
-            genKeys = statementProxy.getTargetStatement().getGeneratedKeys();
+            genKeys = statementProxy.getGeneratedKeys();
         } catch (SQLException e) {
             // java.sql.SQLException: Generated keys not requested. You need to
             // specify Statement.RETURN_GENERATED_KEYS to

diff --git a/rm-datasource/src/test/java/io/seata/rm/datasource/StatementProxyTest.java b/rm-datasource/src/test/java/io/seata/rm/datasource/StatementProxyTest.java
@@ -20,7 +20,11 @@
 import java.sql.Statement;
 import java.sql.Types;
 import java.util.List;
+
+import com.alibaba.druid.mock.MockResultSet;
+import com.alibaba.druid.mock.MockStatement;
 import com.alibaba.druid.pool.DruidDataSource;
+import com.alibaba.druid.util.jdbc.ResultSetMetaDataBase;
 import com.google.common.collect.Lists;
 import io.seata.rm.datasource.mock.MockConnection;
 import io.seata.rm.datasource.mock.MockDriver;
@@ -67,6 +71,10 @@ public static void init() throws SQLException {
 
         Statement statement = mockDriver.createMockStatement((MockConnection)connectionProxy.getTargetConnection());
 
+        MockResultSet mockResultSet = new MockResultSet(statement);
+        ((ResultSetMetaDataBase)mockResultSet.getMetaData()).getColumns().add(new ResultSetMetaDataBase.ColumnMetaData());
+        ((MockStatement) statement).setGeneratedKeys(mockResultSet);
+
         statementProxy = new StatementProxy(connectionProxy, statement);
     }
 

diff --git a/rm-datasource/src/test/java/io/seata/rm/datasource/exec/MySQLInsertExecutorTest.java b/rm-datasource/src/test/java/io/seata/rm/datasource/exec/MySQLInsertExecutorTest.java
@@ -217,9 +217,7 @@ public void testGetPkValuesByAuto_SQLException() {
             columnMetaMap.put(ID_COLUMN, columnMeta);
             when(columnMeta.isAutoincrement()).thenReturn(true);
             when(tableMeta.getPrimaryKeyMap()).thenReturn(columnMetaMap);
-            PreparedStatement preparedStatement = mock(PreparedStatement.class);
-            when(statementProxy.getTargetStatement()).thenReturn(preparedStatement);
-            when(preparedStatement.getGeneratedKeys()).thenThrow(new SQLException());
+            when(statementProxy.getGeneratedKeys()).thenThrow(new SQLException());
             insertExecutor.getPkValuesByAuto();
         });
     }
@@ -235,7 +233,7 @@ public void testGetPkValuesByAuto_SQLException_WarnLog() throws SQLException {
         PreparedStatement preparedStatement = mock(PreparedStatement.class);
         when(statementProxy.getTargetStatement()).thenReturn(preparedStatement);
         SQLException e = new SQLException("test warn log", MySQLInsertExecutor.ERR_SQL_STATE, 1);
-        when(preparedStatement.getGeneratedKeys()).thenThrow(e);
+        when(statementProxy.getGeneratedKeys()).thenThrow(e);
         ResultSet genKeys = mock(ResultSet.class);
         when(statementProxy.getTargetStatement().executeQuery("SELECT LAST_INSERT_ID()")).thenReturn(genKeys);
         Assertions.assertTrue(insertExecutor.getPkValuesByAuto().isEmpty());
@@ -252,7 +250,7 @@ public void testGetPkValuesByAuto_GeneratedKeys_NoResult() throws SQLException {
         PreparedStatement preparedStatement = mock(PreparedStatement.class);
         when(statementProxy.getTargetStatement()).thenReturn(preparedStatement);
         ResultSet resultSet = mock(ResultSet.class);
-        when(preparedStatement.getGeneratedKeys()).thenReturn(resultSet);
+        when(statementProxy.getGeneratedKeys()).thenReturn(resultSet);
         when(resultSet.next()).thenReturn(false);
         when(resultSet.getObject(1)).thenReturn(PK_VALUE);
         List pkValuesByAuto = insertExecutor.getPkValuesByAuto();
@@ -270,7 +268,7 @@ public void testGetPkValuesByAuto_GeneratedKeys_HasResult() throws SQLException
         PreparedStatement preparedStatement = mock(PreparedStatement.class);
         when(statementProxy.getTargetStatement()).thenReturn(preparedStatement);
         ResultSet resultSet = mock(ResultSet.class);
-        when(preparedStatement.getGeneratedKeys()).thenReturn(resultSet);
+        when(statementProxy.getGeneratedKeys()).thenReturn(resultSet);
         when(resultSet.next()).thenReturn(true).thenReturn(false);
         when(resultSet.getObject(1)).thenReturn(PK_VALUE);
         List<Object> pkValues = new ArrayList<>();
@@ -289,7 +287,7 @@ public void testGetPkValuesByAuto_ExecuteQuery_HasResult() throws SQLException {
         when(tableMeta.getPrimaryKeyMap()).thenReturn(columnMetaMap);
         PreparedStatement preparedStatement = mock(PreparedStatement.class);
         when(statementProxy.getTargetStatement()).thenReturn(preparedStatement);
-        when(preparedStatement.getGeneratedKeys()).thenThrow(new SQLException("", MySQLInsertExecutor.ERR_SQL_STATE));
+        when(statementProxy.getGeneratedKeys()).thenThrow(new SQLException("", MySQLInsertExecutor.ERR_SQL_STATE));
         ResultSet resultSet = mock(ResultSet.class);
         when(preparedStatement.executeQuery(anyString())).thenReturn(resultSet);
         when(resultSet.next()).thenReturn(true).thenReturn(false);

diff --git a/rm-datasource/src/test/java/io/seata/rm/datasource/exec/OracleInsertExecutorTest.java b/rm-datasource/src/test/java/io/seata/rm/datasource/exec/OracleInsertExecutorTest.java
@@ -133,9 +133,7 @@ public void testStatement_pkValueByAuto_NotSupportYetException() throws Exceptio
         doReturn(map).when(tableMeta).getPrimaryKeyMap();
 
         ResultSet rs = mock(ResultSet.class);
-        Statement statement = mock(Statement.class);
-        doReturn(statement).when(statementProxy).getTargetStatement();
-        doReturn(rs).when(statement).getGeneratedKeys();
+        doReturn(rs).when(statementProxy).getGeneratedKeys();
         doReturn(false).when(rs).next();
 
         Assertions.assertThrows(NotSupportYetException.class, () -> {
@@ -155,9 +153,7 @@ public void testStatement_pkValueByAuto_NotSupportYetException() throws Exceptio
     public void testGetPkValuesByAuto_NotSupportYetException() {
         Assertions.assertThrows(NotSupportYetException.class, () -> {
             doReturn(tableMeta).when(insertExecutor).getTableMeta();
-            PreparedStatement preparedStatement = mock(PreparedStatement.class);
-            when(statementProxy.getTargetStatement()).thenReturn(preparedStatement);
-            when(preparedStatement.getGeneratedKeys()).thenReturn(mock(ResultSet.class));
+            when(statementProxy.getGeneratedKeys()).thenReturn(mock(ResultSet.class));
             Map<String, ColumnMeta> columnMetaMap = new HashMap<>();
             columnMetaMap.put(ID_COLUMN, new ColumnMeta());
             columnMetaMap.put(USER_ID_COLUMN, new ColumnMeta());