single result set limit error

apache · May 13, 2016 · e5d0dc5 · e5d0dc5
1 parent 23a8815
commit e5d0dc5
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 12 deletions.
diff --git a/...ng-jdbc-core/src/main/java/com/dangdang/ddframe/rdb/sharding/merger/ResultSetFactory.java b/...ng-jdbc-core/src/main/java/com/dangdang/ddframe/rdb/sharding/merger/ResultSetFactory.java
@@ -59,7 +59,7 @@ public static ResultSet getResultSet(final List<ResultSet> resultSets, final Mer
         }
         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 @@ private static ResultSet buildCoupling(final ResultSet preResultSet, final Merge
         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 @@ private static <T> ComponentResultSet<T> join(final ComponentResultSet<T> result
         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;
+        }
+    }
 }
diff --git a/...rc/test/java/com/dangdang/ddframe/rdb/sharding/merger/iterator/IteratorResultSetTest.java b/...rc/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 void assertNextWithLimitForPartData() throws SQLException {
         }
         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));
+    }
 }