Adds support for skips and limits to queryIter

codex-storage · May 14, 2024 · 9dc532b · 9dc532b
1 parent e6ba737
commit 9dc532b
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 9 deletions.
diff --git a/leveldbstatic.nim b/leveldbstatic.nim
@@ -426,25 +426,37 @@ proc getIterValue(iterPtr: ptr leveldb_iterator_t): string =
   str = leveldb_iter_value(iterPtr, addr len)
   return newString(str, len)
 
-proc queryIter*(self: LevelDb, prefix: string = "", keysOnly: bool = false): LevelDbQueryIter = 
-  var iterPtr = leveldb_create_iterator(self.db, self.readOptions)
-
+proc seekToQueryStart(iterPtr: ptr leveldb_iterator_t, prefix: string, skip: int) =
   if prefix.len > 0:
     leveldb_iter_seek(iterPtr, prefix, prefix.len.csize_t)
   else:
     leveldb_iter_seek_to_first(iterPtr)
+  for i in 0..<skip:
+    leveldb_iter_next(iterPtr)
+
+proc closeIter(iter: LevelDbQueryIter, iterPtr: ptr leveldb_iterator_t) =
+  iter.finished = true
+  leveldb_iter_destroy(iterPtr)
+
+proc queryIter*(self: LevelDb, prefix: string = "", keysOnly: bool = false, skip: int = 0, limit: int = 0): LevelDbQueryIter = 
+  var iterPtr = leveldb_create_iterator(self.db, self.readOptions)
 
-  var iter = LevelDbQueryIter()
+  seekToQueryStart(iterPtr, prefix, skip)
+
+  var
+    iter = LevelDbQueryIter()
+    remaining = limit
   let emptyResponse = ("", "")
 
   proc getNext(): (string, string) {.gcsafe, closure.} =
     if iter.finished:
       return emptyResponse
 
-    if leveldb_iter_valid(iterPtr) == levelDbFalse:
-      iter.finished = true
-      leveldb_iter_destroy(iterPtr)
+    if leveldb_iter_valid(iterPtr) == levelDbFalse or (limit > 0 and remaining == 0):
+      iter.closeIter(iterPtr)
       return emptyResponse
+    if limit > 0:
+      dec remaining
 
     let
       keyStr = getIterKey(iterPtr)
@@ -460,8 +472,7 @@ proc queryIter*(self: LevelDb, prefix: string = "", keysOnly: bool = false): Lev
       if keyStr.startsWith(prefix):
         return (keyStr, valueStr)
       else:
-        iter.finished = true
-        leveldb_iter_destroy(iterPtr)
+        iter.closeIter(iterPtr)
         return emptyResponse
     else:
       return (keyStr, valueStr)

diff --git a/tests/test.nim b/tests/test.nim
@@ -241,6 +241,28 @@ suite "leveldb queryIter":
       iter.next() == empty
       iter.finished
 
+  test "skip":
+    let iter = db.queryIter(skip = 1)
+    check:
+      not iter.finished
+      iter.next() == (k2, v2)
+      not iter.finished
+      iter.next() == (k3, v3)
+      not iter.finished
+      iter.next() == empty
+      iter.finished
+
+  test "limit":
+    let iter = db.queryIter(limit = 2)
+    check:
+      not iter.finished
+      iter.next() == (k1, v1)
+      not iter.finished
+      iter.next() == (k2, v2)
+      not iter.finished
+      iter.next() == empty
+      iter.finished
+
   test "iterates only keys":
     let iter = db.queryIter(keysOnly = true)
     check:
@@ -265,6 +287,24 @@ suite "leveldb queryIter":
       iter.next() == empty
       iter.finished
 
+  test "iterates only 'k', skip":
+    let iter = db.queryIter(prefix = "k", skip = 1)
+    check:
+      not iter.finished
+      iter.next() == (k2, v2)
+      not iter.finished
+      iter.next() == empty
+      iter.finished
+
+  test "iterate only 'k', limit":
+    let iter = db.queryIter(prefix = "k", limit = 1)
+    check:
+      not iter.finished
+      iter.next() == (k1, v1)
+      not iter.finished
+      iter.next() == empty
+      iter.finished
+
   test "iterates only 'k', only keys":
     let iter = db.queryIter(prefix = "k", keysOnly = true)
     check: