From 6ab4faa4108a2794bb7443549390975556c044a3 Mon Sep 17 00:00:00 2001
From: Patricio Echague <patricioe@gmail.com>
Date: Mon, 13 Feb 2012 09:52:06 -0800
Subject: [PATCH] Issue #419. Ability to specify start and end key

in KeyIterator
---
 .../cassandra/service/KeyIterator.java        | 30 +++++++--
 .../cassandra/service/KeyIteratorTest.java    | 65 +++++++++++++++++++
 2 files changed, 89 insertions(+), 6 deletions(-)
 create mode 100644 core/src/test/java/me/prettyprint/cassandra/service/KeyIteratorTest.java

diff --git a/core/src/main/java/me/prettyprint/cassandra/service/KeyIterator.java b/core/src/main/java/me/prettyprint/cassandra/service/KeyIterator.java
index b69fc5b75..bf16fad8c 100644
--- a/core/src/main/java/me/prettyprint/cassandra/service/KeyIterator.java
+++ b/core/src/main/java/me/prettyprint/cassandra/service/KeyIterator.java
@@ -24,7 +24,7 @@
 public class KeyIterator<K> implements Iterable<K> {
   private static StringSerializer stringSerializer = new StringSerializer();
 
-  private int maxRowCount = 500;
+  private static int MAX_ROW_COUNT_DEFAULT = 500;
   private int maxColumnCount = 2;	// we only need this to tell if there are any columns in the row (to test for tombstones)
 
   private Iterator<Row<K, String, String>> rowsIterator = null;
@@ -33,6 +33,8 @@ public class KeyIterator<K> implements Iterable<K> {
 
   private K nextValue = null;
   private K lastReadValue = null;
+  private K endKey;
+  private boolean firstRun = true;
 
   private Iterator<K> keyIterator = new Iterator<K>() {
     @Override
@@ -66,22 +68,35 @@ private void findNext(boolean fromRunQuery) {
       }
     }
     if (!rowsIterator.hasNext() && nextValue == null) {
-      runQuery(lastReadValue);
+      runQuery(lastReadValue, endKey);
     }
   }
 
   public KeyIterator(Keyspace keyspace, String columnFamily, AbstractSerializer<K> serializer) {
+    this(keyspace, columnFamily, serializer, null, null, MAX_ROW_COUNT_DEFAULT);
+  }
+
+  public KeyIterator(Keyspace keyspace, String columnFamily, AbstractSerializer<K> serializer, int maxRowCount) {
+    this(keyspace, columnFamily, serializer, null, null, maxRowCount);
+  }
+  
+  public KeyIterator(Keyspace keyspace, String columnFamily, AbstractSerializer<K> serializer, K start, K end) {
+    this(keyspace, columnFamily, serializer, start, end, MAX_ROW_COUNT_DEFAULT);
+  }
+
+  public KeyIterator(Keyspace keyspace, String columnFamily, AbstractSerializer<K> serializer, K start, K end, int maxRowCount) {
     query = HFactory
       .createRangeSlicesQuery(keyspace, serializer, stringSerializer, stringSerializer)
       .setColumnFamily(columnFamily)
       .setRange(null, null, false, maxColumnCount)
       .setRowCount(maxRowCount);
 
-    runQuery(null);
+    endKey = end;
+    runQuery(start, end);
   }
 
-  private void runQuery(K start) {
-    query.setKeys(start, null);
+  private void runQuery(K start, K end) {
+    query.setKeys(start, end);
 
     rowsIterator = null;
     QueryResult<OrderedRows<K, String, String>> result = query.execute();
@@ -89,7 +104,10 @@ private void runQuery(K start) {
     rowsIterator = (rows != null) ? rows.iterator() : null;
 
     // we'll skip this first one, since it is the same as the last one from previous time we executed
-    if (start != null && rowsIterator != null) rowsIterator.next();   
+    if (!firstRun  && rowsIterator != null) 
+      rowsIterator.next();
+
+    firstRun = false;
 
     if (!rowsIterator.hasNext()) {
       nextValue = null;    // all done.  our iterator's hasNext() will now return false;
diff --git a/core/src/test/java/me/prettyprint/cassandra/service/KeyIteratorTest.java b/core/src/test/java/me/prettyprint/cassandra/service/KeyIteratorTest.java
new file mode 100644
index 000000000..5c52d4502
--- /dev/null
+++ b/core/src/test/java/me/prettyprint/cassandra/service/KeyIteratorTest.java
@@ -0,0 +1,65 @@
+package me.prettyprint.cassandra.service;
+
+import static me.prettyprint.hector.api.factory.HFactory.createColumn;
+import static me.prettyprint.hector.api.factory.HFactory.createKeyspace;
+import static me.prettyprint.hector.api.factory.HFactory.createMutator;
+import static me.prettyprint.hector.api.factory.HFactory.getOrCreateCluster;
+import static org.junit.Assert.assertEquals;
+import me.prettyprint.cassandra.BaseEmbededServerSetupTest;
+import me.prettyprint.cassandra.serializers.IntegerSerializer;
+import me.prettyprint.cassandra.serializers.StringSerializer;
+import me.prettyprint.hector.api.Cluster;
+import me.prettyprint.hector.api.Keyspace;
+import me.prettyprint.hector.api.mutation.Mutator;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+public class KeyIteratorTest extends BaseEmbededServerSetupTest {
+
+  private static final StringSerializer se = new StringSerializer();
+  private static final IntegerSerializer is = IntegerSerializer.get();
+  
+  private static final String CF = "Standard1";
+
+  private Cluster cluster;
+  private Keyspace keyspace;
+
+  @Before
+  public void setupCase() {
+    cluster = getOrCreateCluster("Test Cluster", "127.0.0.1:9170");
+    keyspace = createKeyspace("Keyspace1", cluster);
+  }
+
+  @After
+  public void teardownCase() {
+    keyspace = null;
+    cluster = null;
+  }
+
+  @Test
+  public void testIterator() {
+    // Insert 100 rows
+    Mutator<String> m = createMutator(keyspace, se);
+    for (int i = 1; i <= 9; i++) {
+      m.addInsertion("k" + i, CF, createColumn(new Integer(i), new Integer(i), is, is));
+    }
+    m.execute();
+
+    assertKeys(5, "k5", null);
+    assertKeys(9, null, null);
+    assertKeys(7, null, "k7");
+  }
+
+  private void assertKeys(int expected, String start, String end) {
+    Iterable<String> it = new KeyIterator<String>(keyspace, CF, se, start, end);
+
+    int tot = 0;
+    for (String key : it)
+      tot++;
+
+    assertEquals(expected, tot);
+  }
+
+}