1use crate::{other_io_err, DBAndColumns, DBKeyValue};
19use rocksdb::{DBIterator, Direction, IteratorMode, ReadOptions};
20use std::io;
21
22pub trait IterationHandler {
24 type Iterator: Iterator<Item = io::Result<DBKeyValue>>;
25
26 fn iter(self, col: u32, read_opts: ReadOptions) -> Self::Iterator;
30 fn iter_with_prefix(self, col: u32, prefix: &[u8], read_opts: ReadOptions) -> Self::Iterator;
35}
36
37impl<'a> IterationHandler for &'a DBAndColumns {
38 type Iterator = EitherIter<KvdbAdapter<DBIterator<'a>>, std::iter::Once<io::Result<DBKeyValue>>>;
39
40 fn iter(self, col: u32, read_opts: ReadOptions) -> Self::Iterator {
41 match self.cf(col as usize) {
42 Ok(cf) => EitherIter::A(KvdbAdapter(self.db.iterator_cf_opt(cf, read_opts, IteratorMode::Start))),
43 Err(e) => EitherIter::B(std::iter::once(Err(e))),
44 }
45 }
46
47 fn iter_with_prefix(self, col: u32, prefix: &[u8], read_opts: ReadOptions) -> Self::Iterator {
48 match self.cf(col as usize) {
49 Ok(cf) => EitherIter::A(KvdbAdapter(self.db.iterator_cf_opt(
50 cf,
51 read_opts,
52 IteratorMode::From(prefix, Direction::Forward),
53 ))),
54 Err(e) => EitherIter::B(std::iter::once(Err(e))),
55 }
56 }
57}
58
59pub enum EitherIter<A, B> {
61 A(A),
62 B(B),
63}
64
65impl<A, B, I> Iterator for EitherIter<A, B>
66where
67 A: Iterator<Item = I>,
68 B: Iterator<Item = I>,
69{
70 type Item = I;
71
72 fn next(&mut self) -> Option<Self::Item> {
73 match self {
74 Self::A(a) => a.next(),
75 Self::B(b) => b.next(),
76 }
77 }
78}
79
80pub struct KvdbAdapter<T>(T);
82
83impl<T> Iterator for KvdbAdapter<T>
84where
85 T: Iterator<Item = Result<(Box<[u8]>, Box<[u8]>), rocksdb::Error>>,
86{
87 type Item = io::Result<DBKeyValue>;
88
89 fn next(&mut self) -> Option<Self::Item> {
90 self.0
91 .next()
92 .map(|r| r.map_err(other_io_err).map(|(k, v)| (k.into_vec().into(), v.into())))
93 }
94}