Merge pull request #33 from scalableminds/fix-listkeys-omissions

Fix listKeys omissions
scalableminds · Jan 6, 2020 · fe8a923 · fe8a923
2 parents 273a251 + ff79f6a
commit fe8a923
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 9 deletions.
diff --git a/src/main/scala/com/scalableminds/fossildb/db/RocksDBStore.scala b/src/main/scala/com/scalableminds/fossildb/db/RocksDBStore.scala
@@ -112,6 +112,10 @@ class RocksDBKeyIterator(it: RocksIterator, prefix: Option[String]) extends Iter
     key
   }
 
+  def peek: String = {
+    new String(it.key().map(_.toChar))
+  }
+
 }
 
 class RocksDBIterator(it: RocksIterator, prefix: Option[String]) extends Iterator[KeyValuePair[Array[Byte]]] {
@@ -138,7 +142,7 @@ class RocksDBStore(db: RocksDB, handle: ColumnFamilyHandle) {
     new RocksDBIterator(it, prefix)
   }
 
-  def scanKeysOnly(key: String, prefix: Option[String]): Iterator[String] = {
+  def scanKeysOnly(key: String, prefix: Option[String]): RocksDBKeyIterator = {
     val it = db.newIterator(handle)
     it.seek(key.getBytes())
     new RocksDBKeyIterator(it, prefix)

diff --git a/src/main/scala/com/scalableminds/fossildb/db/StoreManager.scala b/src/main/scala/com/scalableminds/fossildb/db/StoreManager.scala
@@ -69,16 +69,12 @@ class StoreManager(dataDir: Path, backupDir: Path, columnFamilies: List[String],
   def compactAllData() = {
     failDuringBackup
     failDuringRestore
-    try {
-      rocksDBManager.get.compactAllData()
-    }
+    rocksDBManager.get.compactAllData()
   }
 
   def exportDB(newDataDir: String, newOptionsFilePathOpt: Option[String]) = {
     failDuringRestore
-    try {
-      rocksDBManager.get.exportToNewDB(Paths.get(newDataDir), newOptionsFilePathOpt)
-    }
+    rocksDBManager.get.exportToNewDB(Paths.get(newDataDir), newOptionsFilePathOpt)
   }
 
   def close = {

diff --git a/src/main/scala/com/scalableminds/fossildb/db/VersionedKeyValueStore.scala b/src/main/scala/com/scalableminds/fossildb/db/VersionedKeyValueStore.scala
@@ -58,6 +58,12 @@ class VersionFilterIterator[T](it: Iterator[KeyValuePair[T]], version: Option[Lo
 
 class KeyOnlyIterator[T](underlying: RocksDBStore, startAfterKey: Option[String]) extends Iterator[String] {
 
+  /*
+     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.
+   */
+
   private var currentKey: Option[String] = startAfterKey
 
   private def compositeKeyFor(keyOpt: Option[String]) = keyOpt match {
@@ -67,13 +73,13 @@ class KeyOnlyIterator[T](underlying: RocksDBStore, startAfterKey: Option[String]
 
   override def hasNext: Boolean = {
     val it = underlying.scanKeysOnly(compositeKeyFor(currentKey), None)
-    if (it.hasNext && currentKey.isDefined) it.next
+    if (it.hasNext && currentKey.isDefined && currentKey.contains(VersionedKey(it.peek).get.key)) it.next
     it.hasNext
   }
 
   override def next(): String = {
     val it = underlying.scanKeysOnly(compositeKeyFor(currentKey), None)
-    if (currentKey.isDefined) it.next
+    if (it.hasNext && currentKey.isDefined && currentKey.contains(VersionedKey(it.peek).get.key)) it.next
     val nextKey = VersionedKey(it.next).get.key
     currentKey = Some(nextKey)
     nextKey

diff --git a/src/test/scala/com/scalableminds/fossildb/FossilDBSuite.scala b/src/test/scala/com/scalableminds/fossildb/FossilDBSuite.scala
@@ -162,6 +162,18 @@ class FossilDBSuite extends FlatSpec with BeforeAndAfterEach with TestHelpers {
     assert(reply2.keys.length == 1)
   }
 
+  it should "return all keys despite lexicographic similarity" in {
+    client.put(PutRequest(collectionA, "abb/1/1-[1,1,1]", Some(1), testData1))
+    client.put(PutRequest(collectionA, "abc/1/1481800838-[3600,2717,121]", Some(123), testData2))
+    client.put(PutRequest(collectionA, "abc/1/1481800839-[3601,2717,121]", Some(123), testData3))
+    client.put(PutRequest(collectionA, "abc/1/1481800839-[3601,2717,121]", Some(125), testData3))
+    client.put(PutRequest(collectionA, "abc/1/1481800839-[3601,2717,121]", Some(128), testData3))
+    client.put(PutRequest(collectionA, "abc/1/1481800846-[3602,2717,121]", Some(123), testData2))
+
+    val reply = client.listKeys(ListKeysRequest(collectionA, None, Some("abb")))
+    assert(reply.keys.length == 3)
+  }
+
   "GetMultipleVersions" should "return all versions in decending order if called without limits" in {
     client.put(PutRequest(collectionA, aKey, Some(0), testData1))
     client.put(PutRequest(collectionA, aKey, Some(1), testData2))