Merge pull request #15 from scalableminds/getMultipleKeys-paginated

Add optional limit for getMultipleKeys

Author: jfrohnhofen

src/main/protobuf/fossildbapi.proto

@@ -64,6 +64,7 @@ message GetMultipleKeysRequest {
     required string key = 2;
     optional string prefix = 3;
     optional uint64 version = 4;
+    optional uint32 limit = 5;
 }
 
 message GetMultipleKeysReply {

src/main/scala/com/scalableminds/fossildb/FossilDBGrpcImpl.scala

@@ -53,7 +53,7 @@ class FossilDBGrpcImpl(storeManager: StoreManager) extends FossilDBGrpc.FossilDB
 
   override def getMultipleKeys(req: GetMultipleKeysRequest) = withExceptionHandler(req) {
     val store = storeManager.getStore(req.collection)
-    val (keys, values) = store.getMultipleKeys(req.key, req.prefix, req.version)
+    val (keys, values) = store.getMultipleKeys(req.key, req.prefix, req.version, req.limit)
     GetMultipleKeysReply(true, None, keys, values.map(ByteString.copyFrom(_)))
   } {errorMsg => GetMultipleKeysReply(false, errorMsg)}
 

src/main/scala/com/scalableminds/fossildb/db/VersionedKeyValueStore.scala

@@ -112,11 +112,11 @@ class VersionedKeyValueStore(underlying: RocksDBStore) {
     }
   }
 
-  def getMultipleKeys(key: String, prefix: Option[String] = None, version: Option[Long] = None): (Seq[String], Seq[Array[Byte]]) = {
+  def getMultipleKeys(key: String, prefix: Option[String] = None, version: Option[Long] = None, limit: Option[Int]): (Seq[String], Seq[Array[Byte]]) = {
     requireValidKey(key)
     prefix.map{p => requireValidKey(p)}
     val iterator: Iterator[VersionedKeyValuePair[Array[Byte]]] = scanKeys(key, prefix, version)
-    val asSequence = iterator.toSeq
+    val asSequence = iterator.take(limit.getOrElse(Int.MaxValue)).toSeq
     val keys = asSequence.map(_.key)
     val values = asSequence.map(_.value)
     (keys, values)

src/test/scala/com/scalableminds/fossildb/FossilDBSuite.scala

@@ -237,6 +237,21 @@ class FossilDBSuite extends FlatSpec with BeforeAndAfterEach {
     assert(reply.values.contains(testData2))
   }
 
+  it should "with limit return only the first n keys of matching version " in {
+    client.put(PutRequest(collectionA, aKey, Some(0), testData1))
+    client.put(PutRequest(collectionA, anotherKey, Some(0), testData1))
+    client.put(PutRequest(collectionA, aThirdKey, Some(0), testData1))
+    client.put(PutRequest(collectionA, aKey, Some(1), testData2))
+    client.put(PutRequest(collectionA, anotherKey, Some(1), testData2))
+    client.put(PutRequest(collectionA, aThirdKey, Some(1), testData2))
+    client.put(PutRequest(collectionA, aKey, Some(2), testData3))
+    client.put(PutRequest(collectionA, anotherKey, Some(2), testData3))
+    client.put(PutRequest(collectionA, aThirdKey, Some(2), testData3))
+    val reply = client.getMultipleKeys(GetMultipleKeysRequest(collectionA, aKey, None, Some(1), Some(2)))
+    assert(reply.keys.length == 2)
+    assert(reply.values.contains(testData2))
+  }
+
   "Backup" should "create non-empty backup directory" in {
     client.put(PutRequest(collectionA, aKey, Some(0), testData1))
     client.backup(BackupRequest())