1  # Nim-RocksDB
  2  # Copyright 2024-2025 Status Research & Development GmbH
  3  # Licensed under either of
  4  #
  5  #  * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE) or http://www.apache.org/licenses/LICENSE-2.0)
  6  #  * GPL license, version 2.0, ([LICENSE-GPLv2](LICENSE-GPLv2) or https://www.gnu.org/licenses/old-licenses/gpl-2.0.en.html)
  7  #
  8  # at your option. This file may not be copied, modified, or distributed except according to those terms.
  9  
 10  ## A `RocksIteratorRef` is a reference to a RocksDB iterator which supports
 11  ## iterating over the key value pairs in a column family.
 12  
 13  {.push raises: [].}
 14  
 15  import ./lib/librocksdb, ./internal/utils, ./options/readopts, ./rocksresult, rocksslice
 16  
 17  export rocksresult, rocksslice
 18  
 19  type
 20    RocksIteratorPtr* = ptr rocksdb_iterator_t
 21  
 22    RocksIteratorRef* = ref object
 23      cPtr: RocksIteratorPtr
 24      readOpts: ReadOptionsRef
 25  
 26    PinnableSlicePtr = ptr rocksdb_pinnableslice_t
 27  
 28    MultiGetIteratorRef* = ref seq[PinnableSlicePtr]
 29  
 30  func init*(T: type MultiGetIteratorRef, len: int): T =
 31    doAssert len > 0
 32    let iter = new T
 33    iter[] =
 34      when NimMajor >= 2 and NimMinor >= 2:
 35        newSeqUninit[PinnableSlicePtr](len)
 36      else:
 37        newSeq[PinnableSlicePtr](len)
 38    iter
 39  
 40  func isClosed*(iter: MultiGetIteratorRef): bool =
 41    iter[].len() == 0
 42  
 43  func close*(iter: MultiGetIteratorRef) =
 44    for pSlice in iter[]:
 45      rocksdb_pinnableslice_destroy(pSlice)
 46    iter[].setLen(0)
 47  
 48  iterator items*(
 49      iter: MultiGetIteratorRef, autoClose: static bool = true
 50  ): Opt[RocksDbSlice] =
 51    ## Iterates over the values in the iterator returning an optional slice
 52    ## for each value. The iterator is automatically closed after the iteration
 53    ## unless autoClose is set to false.
 54    doAssert not iter.isClosed()
 55  
 56    when autoClose:
 57      defer:
 58        iter.close()
 59  
 60    for pSlice in iter[]:
 61      if pSlice.isNil():
 62        yield Opt.none(RocksDbSlice)
 63        continue
 64  
 65      var len: csize_t = 0
 66      let data = rocksdb_pinnableslice_value(pSlice, len.addr)
 67      yield Opt.some(RocksDbSlice.init(data, len))
 68  
 69  proc newRocksIterator*(
 70      cPtr: RocksIteratorPtr, readOpts: ReadOptionsRef
 71  ): RocksIteratorRef =
 72    doAssert not cPtr.isNil()
 73    RocksIteratorRef(cPtr: cPtr, readOpts: readOpts)
 74  
 75  template isClosed*(iter: RocksIteratorRef): bool =
 76    ## Returns `true` if the iterator is closed and `false` otherwise.
 77    iter.cPtr.isNil()
 78  
 79  proc seekToKey*(iter: RocksIteratorRef, key: openArray[byte]) =
 80    ## Seeks to the `key` argument in the column family. If the return code is
 81    ## `false`, the iterator has become invalid and should be closed.
 82    ##
 83    ## It is not clear what happens when the `key` does not exist in the column
 84    ## family. The guess is that the interation will proceed at the next key
 85    ## position. This is suggested by a comment from the GO port at
 86    ##
 87    ##    //github.com/DanielMorsing/rocksdb/blob/master/iterator.go:
 88    ##
 89    ##    Seek moves the iterator the position of the key given or, if the key
 90    ##    doesn't exist, the next key that does exist in the database. If the
 91    ##    key doesn't exist, and there is no next key, the Iterator becomes
 92    ##    invalid.
 93    ##
 94    doAssert not iter.isClosed()
 95    rocksdb_iter_seek(iter.cPtr, cast[cstring](key.unsafeAddrOrNil()), csize_t(key.len))
 96  
 97  proc seekToFirst*(iter: RocksIteratorRef) =
 98    ## Seeks to the first entry in the column family.
 99    doAssert not iter.isClosed()
100    rocksdb_iter_seek_to_first(iter.cPtr)
101  
102  proc seekToLast*(iter: RocksIteratorRef) =
103    ## Seeks to the last entry in the column family.
104    doAssert not iter.isClosed()
105    rocksdb_iter_seek_to_last(iter.cPtr)
106  
107  proc isValid*(iter: RocksIteratorRef): bool =
108    ## Returns `true` if the iterator is valid and `false` otherwise.
109    rocksdb_iter_valid(iter.cPtr).bool
110  
111  proc next*(iter: RocksIteratorRef) =
112    ## Seeks to the next entry in the column family.
113    rocksdb_iter_next(iter.cPtr)
114  
115  proc prev*(iter: RocksIteratorRef) =
116    ## Seeks to the previous entry in the column family.
117    rocksdb_iter_prev(iter.cPtr)
118  
119  template keyOpenArray(iter: RocksIteratorRef): openArray[byte] =
120    let kSlice = rocksdb_iter_key_slice(iter.cPtr)
121    toOpenArray(kSlice.data, kSlice.size)
122  
123  proc key*(iter: RocksIteratorRef, onData: DataProc) =
124    ## Returns the current key using the provided `onData` callback.
125    onData(iter.keyOpenArray())
126  
127  template key*(iter: RocksIteratorRef, asOpenArray: static bool = false): auto =
128    ## Returns the current key.
129    when asOpenArray:
130      iter.keyOpenArray()
131    else:
132      @(iter.keyOpenArray())
133  
134  template valueOpenArray(iter: RocksIteratorRef): openArray[byte] =
135    let vSlice = rocksdb_iter_value_slice(iter.cPtr)
136    toOpenArray(vSlice.data, vSlice.size)
137  
138  proc value*(iter: RocksIteratorRef, onData: DataProc) =
139    ## Returns the current value using the provided `onData` callback.
140    onData(iter.valueOpenArray())
141  
142  template value*(iter: RocksIteratorRef, asOpenArray: static bool = false): auto =
143    ## Returns the current value.
144    when asOpenArray:
145      iter.valueOpenArray()
146    else:
147      @(iter.valueOpenArray())
148  
149  proc status*(iter: RocksIteratorRef): RocksDBResult[void] =
150    ## Returns the status of the iterator.
151    doAssert not iter.isClosed()
152  
153    var errors: cstring
154    rocksdb_iter_get_error(iter.cPtr, cast[cstringArray](errors.addr))
155    bailOnErrors(errors)
156  
157    ok()
158  
159  proc close*(iter: RocksIteratorRef) =
160    ## Closes the `RocksIteratorRef`.
161    if not iter.isClosed():
162      rocksdb_iter_destroy(iter.cPtr)
163      iter.cPtr = nil
164  
165      autoCloseNonNil(iter.readOpts)
166  
167  iterator pairs*(
168      iter: RocksIteratorRef, autoClose: static bool = true
169  ): tuple[key: seq[byte], value: seq[byte]] =
170    ## Iterates over the key value pairs in the column family yielding them in
171    ## the form of a tuple of seq[byte]. The iterator is automatically closed
172    ## after the iteration unless autoClose is set to false.
173    doAssert not iter.isClosed()
174    when autoClose:
175      defer:
176        iter.close()
177  
178    iter.seekToFirst()
179  
180    while iter.isValid():
181      yield (iter.key(), iter.value())
182      iter.next()
183  
184  func keySlice(iter: RocksIteratorRef): RocksDbSlice =
185    ## Returns the current key as a slice.
186    let kSlice = rocksdb_iter_key_slice(iter.cPtr)
187    RocksDbSlice.init(kSlice.data, kSlice.size)
188  
189  func valueSlice(iter: RocksIteratorRef): RocksDbSlice =
190    ## Returns the current value as a slice.
191    let vSlice = rocksdb_iter_value_slice(iter.cPtr)
192    RocksDbSlice.init(vSlice.data, vSlice.size)
193  
194  iterator slicePairs*(
195      iter: RocksIteratorRef, autoClose: static bool = true
196  ): tuple[key: RocksDbSlice, value: RocksDbSlice] =
197    ## Iterates over the key value pairs in the column family yielding them in
198    ## the form of a tuple of slices. The iterator is automatically closed
199    ## after the iteration unless autoClose is set to false.
200    doAssert not iter.isClosed()
201    when autoClose:
202      defer:
203        iter.close()
204  
205    iter.seekToFirst()
206  
207    while iter.isValid():
208      yield (iter.keySlice(), iter.valueSlice())
209      iter.next()