kvdb_memorydb/
lib.rs

1// Copyright 2020 Parity Technologies
2//
3// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
4// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
5// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
6// option. This file may not be copied, modified, or distributed
7// except according to those terms.
8
9use kvdb::{DBKeyValue, DBOp, DBTransaction, DBValue, KeyValueDB};
10use parking_lot::RwLock;
11use std::{
12	collections::{BTreeMap, HashMap},
13	io,
14};
15
16/// A key-value database fulfilling the `KeyValueDB` trait, living in memory.
17/// This is generally intended for tests and is not particularly optimized.
18#[derive(Default)]
19pub struct InMemory {
20	columns: RwLock<HashMap<u32, BTreeMap<Vec<u8>, DBValue>>>,
21}
22
23/// Create an in-memory database with the given number of columns.
24/// Columns will be indexable by 0..`num_cols`
25pub fn create(num_cols: u32) -> InMemory {
26	let mut cols = HashMap::new();
27
28	for idx in 0..num_cols {
29		cols.insert(idx, BTreeMap::new());
30	}
31
32	InMemory { columns: RwLock::new(cols) }
33}
34
35fn invalid_column(col: u32) -> io::Error {
36	io::Error::new(io::ErrorKind::Other, format!("No such column family: {:?}", col))
37}
38
39impl KeyValueDB for InMemory {
40	fn get(&self, col: u32, key: &[u8]) -> io::Result<Option<DBValue>> {
41		let columns = self.columns.read();
42		match columns.get(&col) {
43			None => Err(invalid_column(col)),
44			Some(map) => Ok(map.get(key).cloned()),
45		}
46	}
47
48	fn get_by_prefix(&self, col: u32, prefix: &[u8]) -> io::Result<Option<DBValue>> {
49		let columns = self.columns.read();
50		match columns.get(&col) {
51			None => Err(invalid_column(col)),
52			Some(map) => Ok(map.iter().find(|&(ref k, _)| k.starts_with(prefix)).map(|(_, v)| v.to_vec())),
53		}
54	}
55
56	fn write(&self, transaction: DBTransaction) -> io::Result<()> {
57		let mut columns = self.columns.write();
58		let ops = transaction.ops;
59		for op in ops {
60			match op {
61				DBOp::Insert { col, key, value } =>
62					if let Some(col) = columns.get_mut(&col) {
63						col.insert(key.into_vec(), value);
64					},
65				DBOp::Delete { col, key } =>
66					if let Some(col) = columns.get_mut(&col) {
67						col.remove(&*key);
68					},
69				DBOp::DeletePrefix { col, prefix } =>
70					if let Some(col) = columns.get_mut(&col) {
71						use std::ops::Bound;
72						if prefix.is_empty() {
73							col.clear();
74						} else {
75							let start_range = Bound::Included(prefix.to_vec());
76							let keys: Vec<_> = if let Some(end_range) = kvdb::end_prefix(&prefix[..]) {
77								col.range((start_range, Bound::Excluded(end_range)))
78									.map(|(k, _)| k.clone())
79									.collect()
80							} else {
81								col.range((start_range, Bound::Unbounded)).map(|(k, _)| k.clone()).collect()
82							};
83							for key in keys.into_iter() {
84								col.remove(&key[..]);
85							}
86						}
87					},
88			}
89		}
90		Ok(())
91	}
92
93	fn iter<'a>(&'a self, col: u32) -> Box<dyn Iterator<Item = io::Result<DBKeyValue>> + 'a> {
94		match self.columns.read().get(&col) {
95			Some(map) => Box::new(
96				// TODO: worth optimizing at all?
97				map.clone().into_iter().map(|(k, v)| Ok((k.into(), v))),
98			),
99			None => Box::new(std::iter::once(Err(invalid_column(col)))),
100		}
101	}
102
103	fn iter_with_prefix<'a>(
104		&'a self,
105		col: u32,
106		prefix: &'a [u8],
107	) -> Box<dyn Iterator<Item = io::Result<DBKeyValue>> + 'a> {
108		match self.columns.read().get(&col) {
109			Some(map) => Box::new(
110				map.clone()
111					.into_iter()
112					.filter(move |&(ref k, _)| k.starts_with(prefix))
113					.map(|(k, v)| Ok((k.into(), v))),
114			),
115			None => Box::new(std::iter::once(Err(invalid_column(col)))),
116		}
117	}
118}
119
120#[cfg(test)]
121mod tests {
122	use super::create;
123	use kvdb_shared_tests as st;
124	use std::io;
125
126	#[test]
127	fn get_fails_with_non_existing_column() -> io::Result<()> {
128		let db = create(1);
129		st::test_get_fails_with_non_existing_column(&db)
130	}
131
132	#[test]
133	fn put_and_get() -> io::Result<()> {
134		let db = create(1);
135		st::test_put_and_get(&db)
136	}
137
138	#[test]
139	fn delete_and_get() -> io::Result<()> {
140		let db = create(1);
141		st::test_delete_and_get(&db)
142	}
143
144	#[test]
145	fn delete_prefix() -> io::Result<()> {
146		let db = create(st::DELETE_PREFIX_NUM_COLUMNS);
147		st::test_delete_prefix(&db)
148	}
149
150	#[test]
151	fn iter() -> io::Result<()> {
152		let db = create(1);
153		st::test_iter(&db)
154	}
155
156	#[test]
157	fn iter_with_prefix() -> io::Result<()> {
158		let db = create(1);
159		st::test_iter_with_prefix(&db)
160	}
161
162	#[test]
163	fn complex() -> io::Result<()> {
164		let db = create(1);
165		st::test_complex(&db)
166	}
167}
kvdb_memorydb/lib.rs

kvdb_memorydb/
lib.rs
