single result set limit error

sharding-jdbc-core/src/main/java/com/dangdang/ddframe/rdb/sharding/merger/ResultSetFactory.java

@@ -59,7 +59,7 @@ public final class ResultSetFactory {
         }
         if (1 == filteredResultSets.size()) {
             log.trace("Sharding-JDBC: Only one result set");
-            return filteredResultSets.get(0);
+            return joinLimit(filteredResultSets.get(0), mergeContext);
         }
         mergeContext.buildContextWithResultSet((WrapperResultSet) filteredResultSets.get(0));
         return buildCoupling(buildReducer(filteredResultSets, mergeContext), mergeContext);
@@ -97,9 +97,7 @@ public final class ResultSetFactory {
         if (mergeContext.needToSort()) {
             currentResultSet = join(new MemoryOrderByCouplingResultSet(mergeContext.getCurrentOrderByKeys()), currentResultSet);
         }
-        if (mergeContext.hasLimit()) {
-            currentResultSet = join(new LimitCouplingResultSet(mergeContext.getLimit()), currentResultSet);
-        }
+        currentResultSet = joinLimit(currentResultSet, mergeContext);
         return currentResultSet;
     }
     
@@ -108,4 +106,12 @@ public final class ResultSetFactory {
         resultSet.init(preResultSet);
         return resultSet;
     }
+    
+    private static ResultSet joinLimit(final ResultSet preResultSet, final MergeContext mergeContext) throws SQLException {
+        if (mergeContext.hasLimit()) {
+            return join(new LimitCouplingResultSet(mergeContext.getLimit()), preResultSet);
+        } else {
+            return preResultSet;
+        }
+    }
 }

sharding-jdbc-core/src/test/java/com/dangdang/ddframe/rdb/sharding/merger/iterator/IteratorResultSetTest.java

@@ -17,19 +17,19 @@
 
 package com.dangdang.ddframe.rdb.sharding.merger.iterator;
 
-import static org.hamcrest.CoreMatchers.is;
-import static org.junit.Assert.assertThat;
+import com.dangdang.ddframe.rdb.sharding.merger.ResultSetFactory;
+import com.dangdang.ddframe.rdb.sharding.merger.fixture.MockResultSet;
+import com.dangdang.ddframe.rdb.sharding.parser.result.merger.Limit;
+import com.dangdang.ddframe.rdb.sharding.parser.result.merger.MergeContext;
+import org.junit.Test;
 
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.util.Arrays;
+import java.util.Collections;
 
-import org.junit.Test;
-
-import com.dangdang.ddframe.rdb.sharding.merger.ResultSetFactory;
-import com.dangdang.ddframe.rdb.sharding.merger.fixture.MockResultSet;
-import com.dangdang.ddframe.rdb.sharding.parser.result.merger.Limit;
-import com.dangdang.ddframe.rdb.sharding.parser.result.merger.MergeContext;
+import static org.hamcrest.CoreMatchers.is;
+import static org.junit.Assert.assertThat;
 
 public final class IteratorResultSetTest {
     
@@ -66,4 +66,30 @@ public final class IteratorResultSetTest {
         }
         assertThat(count, is(1));
     }
+    
+    @Test
+    public void assertNextWithLimitForSingleResultSet() throws SQLException {
+        MergeContext mergeContext = new MergeContext();
+        mergeContext.setLimit(new Limit(2, 3));
+        ResultSet resultSet = ResultSetFactory.getResultSet(Collections.<ResultSet>singletonList(new MockResultSet<>(1, 2, 3, 4, 5, 6)), mergeContext);
+        int count = 0;
+        while (resultSet.next()) {
+            System.out.println(resultSet.getInt(1));
+            count++;
+        }
+        assertThat(count, is(3));
+    }
+    
+    @Test
+    public void assertNextWithLimitForSingleResultSetNotFull() throws SQLException {
+        MergeContext mergeContext = new MergeContext();
+        mergeContext.setLimit(new Limit(2, 10));
+        ResultSet resultSet = ResultSetFactory.getResultSet(Collections.<ResultSet>singletonList(new MockResultSet<>(1, 2, 3, 4, 5, 6)), mergeContext);
+        int count = 0;
+        while (resultSet.next()) {
+            count++;
+            System.out.println(resultSet.getInt(1));
+        }
+        assertThat(count, is(4));
+    }
 }