use crate::{other_io_err, DBAndColumns, DBKeyValue};
use rocksdb::{DBIterator, Direction, IteratorMode, ReadOptions};
use std::io;
pub trait IterationHandler {
type Iterator: Iterator<Item = io::Result<DBKeyValue>>;
fn iter(self, col: u32, read_opts: ReadOptions) -> Self::Iterator;
fn iter_with_prefix(self, col: u32, prefix: &[u8], read_opts: ReadOptions) -> Self::Iterator;
}
impl<'a> IterationHandler for &'a DBAndColumns {
type Iterator = EitherIter<KvdbAdapter<DBIterator<'a>>, std::iter::Once<io::Result<DBKeyValue>>>;
fn iter(self, col: u32, read_opts: ReadOptions) -> Self::Iterator {
match self.cf(col as usize) {
Ok(cf) => EitherIter::A(KvdbAdapter(self.db.iterator_cf_opt(cf, read_opts, IteratorMode::Start))),
Err(e) => EitherIter::B(std::iter::once(Err(e))),
}
}
fn iter_with_prefix(self, col: u32, prefix: &[u8], read_opts: ReadOptions) -> Self::Iterator {
match self.cf(col as usize) {
Ok(cf) => EitherIter::A(KvdbAdapter(self.db.iterator_cf_opt(
cf,
read_opts,
IteratorMode::From(prefix, Direction::Forward),
))),
Err(e) => EitherIter::B(std::iter::once(Err(e))),
}
}
}
pub enum EitherIter<A, B> {
A(A),
B(B),
}
impl<A, B, I> Iterator for EitherIter<A, B>
where
A: Iterator<Item = I>,
B: Iterator<Item = I>,
{
type Item = I;
fn next(&mut self) -> Option<Self::Item> {
match self {
Self::A(a) => a.next(),
Self::B(b) => b.next(),
}
}
}
pub struct KvdbAdapter<T>(T);
impl<T> Iterator for KvdbAdapter<T>
where
T: Iterator<Item = Result<(Box<[u8]>, Box<[u8]>), rocksdb::Error>>,
{
type Item = io::Result<DBKeyValue>;
fn next(&mut self) -> Option<Self::Item> {
self.0
.next()
.map(|r| r.map_err(other_io_err).map(|(k, v)| (k.into_vec().into(), v.into())))
}
}