From fdeac9bc19cf9e5c44ff62e9ac1a75796993470c Mon Sep 17 00:00:00 2001
From: terrymanu <terrymanu@163.com>
Date: Tue, 12 Sep 2017 15:12:13 +0800
Subject: [PATCH] fix #370

---
 RELEASE-NOTES.md                              |  1 +
 .../statement/ShardingPreparedStatement.java  | 65 +++++++++----------
 2 files changed, 32 insertions(+), 34 deletions(-)

diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md
index fb7cc4b7043b7..36fddbd6b48c5 100644
--- a/RELEASE-NOTES.md
+++ b/RELEASE-NOTES.md
@@ -4,6 +4,7 @@
 
 1. [ISSUE #356](https://github.com/dangdangdotcom/sharding-jdbc/issues/356) 在SQL的Where条件中兼容不是分片列的REGEXP操作符
 1. [ISSUE #362](https://github.com/dangdangdotcom/sharding-jdbc/issues/362) 读写分离使用PreparedStatement并未调用setParameter方法导致出错
+1. [ISSUE #370](https://github.com/dangdangdotcom/sharding-jdbc/issues/370) 使用原生自增主键调用getGeneratedKeys出错
 
 ## 1.5.3
 
diff --git a/sharding-jdbc-core/src/main/java/com/dangdang/ddframe/rdb/sharding/jdbc/core/statement/ShardingPreparedStatement.java b/sharding-jdbc-core/src/main/java/com/dangdang/ddframe/rdb/sharding/jdbc/core/statement/ShardingPreparedStatement.java
index 2aba88601dc8a..2d93cf1fbcf6a 100644
--- a/sharding-jdbc-core/src/main/java/com/dangdang/ddframe/rdb/sharding/jdbc/core/statement/ShardingPreparedStatement.java
+++ b/sharding-jdbc-core/src/main/java/com/dangdang/ddframe/rdb/sharding/jdbc/core/statement/ShardingPreparedStatement.java
@@ -174,19 +174,9 @@ private Collection<PreparedStatement> generatePreparedStatementForDDL(final SQLE
     }
     
     private PreparedStatement generatePreparedStatement(final SQLExecutionUnit sqlExecutionUnit) throws SQLException {
-        Optional<GeneratedKey> generatedKey = getGeneratedKey();
         Connection connection = getConnection().getConnection(sqlExecutionUnit.getDataSource(), routeResult.getSqlStatement().getType());
-        if (returnGeneratedKeys && generatedKey.isPresent()) {
-            return connection.prepareStatement(sqlExecutionUnit.getSql(), RETURN_GENERATED_KEYS);
-        }
-        return connection.prepareStatement(sqlExecutionUnit.getSql(), resultSetType, resultSetConcurrency, resultSetHoldability);
-    }
-    
-    private Optional<GeneratedKey> getGeneratedKey() {
-        if (null != routeResult && routeResult.getSqlStatement() instanceof InsertStatement) {
-            return Optional.fromNullable(((InsertStatement) routeResult.getSqlStatement()).getGeneratedKey());
-        }
-        return Optional.absent();
+        return returnGeneratedKeys ? connection.prepareStatement(sqlExecutionUnit.getSql(), RETURN_GENERATED_KEYS)
+                : connection.prepareStatement(sqlExecutionUnit.getSql(), resultSetType, resultSetConcurrency, resultSetHoldability);
     }
     
     @Override
@@ -211,28 +201,6 @@ public void addBatch() throws SQLException {
         }
     }
     
-    @Override
-    public int[] executeBatch() throws SQLException {
-        try {
-            return new BatchPreparedStatementExecutor(getConnection().getShardingContext().getExecutorEngine(), 
-                    getConnection().getShardingContext().getDatabaseType(), routeResult.getSqlStatement().getType(), batchStatementUnits, parameterSets).executeBatch();
-        } finally {
-            clearBatch();
-        }
-    }
-    
-    @Override
-    public ResultSet getGeneratedKeys() throws SQLException {
-        Optional<GeneratedKey> generatedKey = getGeneratedKey();
-        if (returnGeneratedKeys && generatedKey.isPresent()) {
-            return new GeneratedKeysResultSet(routeResult.getGeneratedKeys().iterator(), generatedKey.get().getColumn(), this);
-        }
-        if (1 == routedStatements.size()) {
-            return routedStatements.iterator().next().getGeneratedKeys();
-        }
-        return new GeneratedKeysResultSet();
-    }
-    
     private List<BatchPreparedStatementUnit> routeBatch() throws SQLException {
         List<BatchPreparedStatementUnit> result = new ArrayList<>();
         routeResult = routingEngine.route(getParameters());
@@ -260,6 +228,35 @@ public boolean apply(final BatchPreparedStatementUnit input) {
         return result;
     }
     
+    @Override
+    public int[] executeBatch() throws SQLException {
+        try {
+            return new BatchPreparedStatementExecutor(getConnection().getShardingContext().getExecutorEngine(), 
+                    getConnection().getShardingContext().getDatabaseType(), routeResult.getSqlStatement().getType(), batchStatementUnits, parameterSets).executeBatch();
+        } finally {
+            clearBatch();
+        }
+    }
+    
+    @Override
+    public ResultSet getGeneratedKeys() throws SQLException {
+        Optional<GeneratedKey> generatedKey = getGeneratedKey();
+        if (returnGeneratedKeys && generatedKey.isPresent()) {
+            return new GeneratedKeysResultSet(routeResult.getGeneratedKeys().iterator(), generatedKey.get().getColumn(), this);
+        }
+        if (1 == routedStatements.size()) {
+            return routedStatements.iterator().next().getGeneratedKeys();
+        }
+        return new GeneratedKeysResultSet();
+    }
+    
+    private Optional<GeneratedKey> getGeneratedKey() {
+        if (null != routeResult && routeResult.getSqlStatement() instanceof InsertStatement) {
+            return Optional.fromNullable(((InsertStatement) routeResult.getSqlStatement()).getGeneratedKey());
+        }
+        return Optional.absent();
+    }
+    
     @Override
     public ResultSet getResultSet() throws SQLException {
         if (null != currentResultSet) {