Fix pagination for GetMultipleKeys

Author: fm3

src/main/protobuf/fossildbapi.proto

@@ -61,7 +61,7 @@ message GetMultipleVersionsReply {
 
 message GetMultipleKeysRequest {
     required string collection = 1;
-    required string key = 2;
+    optional string startAfterKey = 2;
     optional string prefix = 3;
     optional uint64 version = 4;
     optional uint32 limit = 5;

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

@@ -51,7 +51,7 @@ class FossilDBGrpcImpl(storeManager: StoreManager)
 
   override def getMultipleKeys(req: GetMultipleKeysRequest): Future[GetMultipleKeysReply] = withExceptionHandler(req) {
     val store = storeManager.getStore(req.collection)
-    val (keys, values, versions) = store.getMultipleKeys(req.key, req.prefix, req.version, req.limit)
+    val (keys, values, versions) = store.getMultipleKeys(req.startAfterKey, req.prefix, req.version, req.limit)
     GetMultipleKeysReply(success = true, None, keys, values.map(ByteString.copyFrom), versions)
   } { errorMsg => GetMultipleKeysReply(success = false, errorMsg) }
 

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

@@ -5,13 +5,15 @@ import scala.util.Try
 
 
 case class VersionedKey(key: String, version: Long) {
-  override def toString: String = s"$key@${(~version).toHexString.toUpperCase}@$version"
+  override def toString: String = s"$key${VersionedKey.versionSeparator}${(~version).toHexString.toUpperCase}${VersionedKey.versionSeparator}$version"
 }
 
 object VersionedKey {
 
+  val versionSeparator: Char = '@'
+
   def apply(key: String): Option[VersionedKey] = {
-    val parts = key.split('@')
+    val parts = key.split(versionSeparator)
     for {
       key <- parts.headOption
       versionString <- parts.lastOption
@@ -91,8 +93,6 @@ class VersionedKeyValueStore(underlying: RocksDBStore) {
   def get(key: String, version: Option[Long] = None): Option[VersionedKeyValuePair[Array[Byte]]] =
     scanVersionValuePairs(key, version).toStream.headOption
 
-  def getUnderlying: RocksDBStore = underlying
-
   def getMultipleVersions(key: String, oldestVersion: Option[Long] = None, newestVersion: Option[Long] = None): (List[Array[Byte]], List[Long]) = {
 
     @tailrec
@@ -113,33 +113,54 @@ class VersionedKeyValueStore(underlying: RocksDBStore) {
 
   private def scanVersionValuePairs(key: String, version: Option[Long] = None): Iterator[VersionedKeyValuePair[Array[Byte]]] = {
     requireValidKey(key)
-    val prefix = s"$key@"
+    val prefix = s"$key${VersionedKey.versionSeparator}"
     underlying.scan(version.map(VersionedKey(key, _).toString).getOrElse(prefix), Some(prefix)).flatMap { pair =>
       VersionedKey(pair.key).map(VersionedKeyValuePair(_, pair.value))
     }
   }
 
   private def scanVersionsOnly(key: String, version: Option[Long] = None): Iterator[VersionedKey] = {
     requireValidKey(key)
-    val prefix = s"$key@"
+    val prefix = s"$key${VersionedKey.versionSeparator}"
     underlying.scanKeysOnly(version.map(VersionedKey(key, _).toString).getOrElse(prefix), Some(prefix)).flatMap { key =>
       VersionedKey(key)
     }
   }
 
-  def getMultipleKeys(key: String, prefix: Option[String] = None, version: Option[Long] = None, limit: Option[Int]): (Seq[String], Seq[Array[Byte]], Seq[Long]) = {
-    requireValidKey(key)
+  def getMultipleKeys(startAfterKey: Option[String], prefix: Option[String] = None, version: Option[Long] = None, limit: Option[Int]): (Seq[String], Seq[Array[Byte]], Seq[Long]) = {
+    startAfterKey.foreach(requireValidKey)
     prefix.foreach{ p => requireValidKey(p)}
-    val iterator: VersionFilterIterator = scanKeys(key, prefix, version)
-    val asSequence = iterator.take(limit.getOrElse(Int.MaxValue)).toSeq
-    val keys = asSequence.map(_.key)
-    val values = asSequence.map(_.value)
-    val versions = asSequence.map(_.version)
+    val iterator: VersionFilterIterator = scanKeys(startAfterKey, prefix, version)
+
+    /*
+       Note that seek in the underlying iterators either hits precisely or goes to the
+       lexicographically *next* key. To achieve correct behavior with startAfterKey,
+       we have to advance once in case of the exact hit.
+     */
+    val firstItemOpt: Option[VersionedKeyValuePair[Array[Byte]]] = if (iterator.hasNext) {
+      val firstItem = iterator.next()
+      if (startAfterKey.contains(firstItem.key)) {
+        None
+      } else {
+        Some(firstItem)
+      }
+    } else None
+
+    val limitPadded = limit.map(_ - 1).getOrElse(Int.MaxValue)
+    val asVector = iterator.take(limitPadded).toVector
+    val asSequenceAdvancedIfNeeded = firstItemOpt.map(_ +: asVector).getOrElse(asVector).take(limit.getOrElse(Int.MaxValue))
+    val keys = asSequenceAdvancedIfNeeded.map(_.key)
+    println(s"Returning keys ${keys.toList}")
+    val values = asSequenceAdvancedIfNeeded.map(_.value)
+    val versions = asSequenceAdvancedIfNeeded.map(_.version)
     (keys, values, versions)
   }
 
-  private def scanKeys(key: String, prefix: Option[String] = None, version: Option[Long] = None): VersionFilterIterator =
-    new VersionFilterIterator(underlying.scan(key, prefix), version)
+  private def scanKeys(startAfterKey: Option[String], prefix: Option[String] = None, version: Option[Long] = None): VersionFilterIterator = {
+    val fullKey = startAfterKey.map(key => s"$key${VersionedKey.versionSeparator}").getOrElse("")
+    println(s"Scanning to $fullKey")
+    new VersionFilterIterator(underlying.scan(fullKey, prefix), version)
+  }
 
   def deleteMultipleVersions(key: String, oldestVersion: Option[Long] = None, newestVersion: Option[Long] = None): Unit = {
     @tailrec
@@ -178,6 +199,6 @@ class VersionedKeyValueStore(underlying: RocksDBStore) {
   }
 
   private def requireValidKey(key: String): Unit = {
-    require(!(key contains "@"), "keys cannot contain the char @")
+    require(!key.contains(VersionedKey.versionSeparator), s"keys cannot contain the char ${VersionedKey.versionSeparator}")
   }
 }

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

@@ -206,7 +206,7 @@ class FossilDBSuite extends FlatSpec with BeforeAndAfterEach with TestHelpers wi
     client.put(PutRequest(collectionA, aKey, Some(0), testData1))
     client.put(PutRequest(collectionA, anotherKey, Some(0), testData2))
     client.put(PutRequest(collectionA, aThirdKey, Some(0), testData3))
-    val reply = client.getMultipleKeys(GetMultipleKeysRequest(collectionA, aThirdKey))
+    val reply = client.getMultipleKeys(GetMultipleKeysRequest(collectionA, Some(aKey)))
     assert(reply.keys.length == 2)
     assert(reply.keys.contains(anotherKey))
     assert(reply.keys.contains(aThirdKey))
@@ -227,8 +227,8 @@ class FossilDBSuite extends FlatSpec with BeforeAndAfterEach with TestHelpers wi
     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)))
-    assert(reply.keys.length == 3)
+    val reply = client.getMultipleKeys(GetMultipleKeysRequest(collectionA, Some(aKey), None, Some(1)))
+    assert(reply.keys.length == 2)
     assert(reply.values.contains(testData2))
   }
 
@@ -242,7 +242,7 @@ class FossilDBSuite extends FlatSpec with BeforeAndAfterEach with TestHelpers wi
     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, Some("aK"), Some(1)))
+    val reply = client.getMultipleKeys(GetMultipleKeysRequest(collectionA, None, Some("aK"), Some(1)))
     assert(reply.keys.length == 1)
     assert(reply.values.contains(testData2))
     assert(reply.actualVersions.contains(1))
@@ -258,7 +258,7 @@ class FossilDBSuite extends FlatSpec with BeforeAndAfterEach with TestHelpers wi
     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)))
+    val reply = client.getMultipleKeys(GetMultipleKeysRequest(collectionA, None, None, Some(1), Some(2)))
     assert(reply.keys.length == 2)
     assert(reply.values.contains(testData2))
     assert(reply.actualVersions.contains(1))