ckb_rocksdb/
db.rs

1// Copyright 2014 Tyler Neely
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14//
15
16use crate::ffi;
17use crate::ffi_util::to_cpath;
18
19use crate::{
20    db_options::OptionsMustOutliveDB,
21    handle::Handle,
22    open_raw::{OpenRaw, OpenRawFFI},
23    ops,
24    ops::*,
25    ColumnFamily, DBRawIterator, Error, Options, ReadOptions, Snapshot,
26};
27
28use std::collections::BTreeMap;
29use std::ffi::CStr;
30use std::fmt;
31use std::marker::PhantomData;
32use std::path::{Path, PathBuf};
33use std::slice;
34
35/// A RocksDB database.
36///
37/// See crate level documentation for a simple usage example.
38pub struct DB {
39    pub(crate) inner: *mut ffi::rocksdb_t,
40    cfs: BTreeMap<String, ColumnFamily>,
41    path: PathBuf,
42    _outlive: Vec<OptionsMustOutliveDB>,
43}
44
45impl Handle<ffi::rocksdb_t> for DB {
46    fn handle(&self) -> *mut ffi::rocksdb_t {
47        self.inner
48    }
49}
50
51impl ops::Open for DB {}
52impl ops::OpenCF for DB {}
53
54impl OpenRaw for DB {
55    type Pointer = ffi::rocksdb_t;
56    type Descriptor = ();
57
58    fn open_ffi(input: OpenRawFFI<'_, Self::Descriptor>) -> Result<*mut Self::Pointer, Error> {
59        let pointer = unsafe {
60            if input.num_column_families <= 0 {
61                ffi_try!(ffi::rocksdb_open(input.options, input.path,))
62            } else {
63                ffi_try!(ffi::rocksdb_open_column_families(
64                    input.options,
65                    input.path,
66                    input.num_column_families,
67                    input.column_family_names,
68                    input.column_family_options,
69                    input.column_family_handles,
70                ))
71            }
72        };
73
74        Ok(pointer)
75    }
76
77    fn build<I>(
78        path: PathBuf,
79        _open_descriptor: Self::Descriptor,
80        pointer: *mut Self::Pointer,
81        column_families: I,
82        outlive: Vec<OptionsMustOutliveDB>,
83    ) -> Result<Self, Error>
84    where
85        I: IntoIterator<Item = (String, *mut ffi::rocksdb_column_family_handle_t)>,
86    {
87        let cfs: BTreeMap<_, _> = column_families
88            .into_iter()
89            .map(|(k, h)| (k, ColumnFamily::new(h)))
90            .collect();
91
92        Ok(DB {
93            inner: pointer,
94            cfs,
95            path,
96            _outlive: outlive,
97        })
98    }
99}
100
101impl ops::Read for DB {}
102impl ops::Write for DB {}
103
104unsafe impl Send for DB {}
105unsafe impl Sync for DB {}
106
107impl DB {
108    pub fn list_cf<P: AsRef<Path>>(opts: &Options, path: P) -> Result<Vec<String>, Error> {
109        let cpath = to_cpath(
110            path,
111            "Failed to convert path to CString when opening database.",
112        )?;
113        let mut length = 0;
114
115        unsafe {
116            let ptr = ffi_try!(ffi::rocksdb_list_column_families(
117                opts.inner,
118                cpath.as_ptr() as *const _,
119                &mut length,
120            ));
121
122            let vec = slice::from_raw_parts(ptr, length)
123                .iter()
124                .map(|ptr| CStr::from_ptr(*ptr).to_string_lossy().into_owned())
125                .collect();
126            ffi::rocksdb_list_column_families_destroy(ptr, length);
127            Ok(vec)
128        }
129    }
130
131    pub fn destroy<P: AsRef<Path>>(opts: &Options, path: P) -> Result<(), Error> {
132        let cpath = to_cpath(
133            path,
134            "Failed to convert path to CString when opening database.",
135        )?;
136        unsafe {
137            ffi_try!(ffi::rocksdb_destroy_db(opts.inner, cpath.as_ptr(),));
138        }
139        Ok(())
140    }
141
142    pub fn repair<P: AsRef<Path>>(opts: Options, path: P) -> Result<(), Error> {
143        let cpath = to_cpath(
144            path,
145            "Failed to convert path to CString when opening database.",
146        )?;
147        unsafe {
148            ffi_try!(ffi::rocksdb_repair_db(opts.inner, cpath.as_ptr(),));
149        }
150        Ok(())
151    }
152
153    pub fn path(&self) -> &Path {
154        self.path.as_path()
155    }
156
157    pub fn snapshot(&self) -> Snapshot<'_> {
158        let snapshot = unsafe { ffi::rocksdb_create_snapshot(self.inner) };
159        Snapshot {
160            db: self,
161            inner: snapshot,
162        }
163    }
164}
165
166impl Drop for DB {
167    fn drop(&mut self) {
168        unsafe {
169            for cf in self.cfs.values() {
170                ffi::rocksdb_column_family_handle_destroy(cf.inner);
171            }
172            ffi::rocksdb_close(self.inner);
173        }
174    }
175}
176
177impl fmt::Debug for DB {
178    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
179        write!(f, "RocksDB {{ path: {:?} }}", self.path())
180    }
181}
182
183impl Iterate for DB {
184    fn get_raw_iter<'a: 'b, 'b>(&'a self, readopts: &ReadOptions) -> DBRawIterator<'b> {
185        unsafe {
186            DBRawIterator {
187                inner: ffi::rocksdb_create_iterator(self.inner, readopts.handle()),
188                db: PhantomData,
189            }
190        }
191    }
192}
193
194impl IterateCF for DB {
195    fn get_raw_iter_cf<'a: 'b, 'b>(
196        &'a self,
197        cf_handle: &ColumnFamily,
198        readopts: &ReadOptions,
199    ) -> Result<DBRawIterator<'b>, Error> {
200        unsafe {
201            Ok(DBRawIterator {
202                inner: ffi::rocksdb_create_iterator_cf(
203                    self.inner,
204                    readopts.handle(),
205                    cf_handle.inner,
206                ),
207                db: PhantomData,
208            })
209        }
210    }
211}
212
213impl GetColumnFamilys for DB {
214    fn get_cfs(&self) -> &BTreeMap<String, ColumnFamily> {
215        &self.cfs
216    }
217    fn get_mut_cfs(&mut self) -> &mut BTreeMap<String, ColumnFamily> {
218        &mut self.cfs
219    }
220}
221
222#[test]
223fn test_db_vector() {
224    use crate::prelude::*;
225    use libc::size_t;
226    use std::mem;
227
228    let len: size_t = 4;
229    let data = unsafe { libc::calloc(len, mem::size_of::<u8>()) as *mut u8 };
230    let v = unsafe { DBVector::from_c(data, len) };
231    let ctrl = [0u8, 0, 0, 0];
232    assert_eq!(&*v, &ctrl[..]);
233}
234
235#[test]
236fn external() {
237    use crate::{prelude::*, TemporaryDBPath};
238
239    let path = TemporaryDBPath::new();
240    {
241        let db = DB::open_default(&path).unwrap();
242        let p = db.put(b"k1", b"v1111");
243        assert!(p.is_ok());
244        let r: Result<Option<DBVector>, Error> = db.get(b"k1");
245        assert!(r.unwrap().unwrap().to_utf8().unwrap() == "v1111");
246        assert!(db.delete(b"k1").is_ok());
247        assert!(db.get(b"k1").unwrap().is_none());
248    }
249}
250
251#[test]
252fn errors_do_stuff() {
253    use crate::{prelude::*, TemporaryDBPath};
254
255    let path = TemporaryDBPath::new();
256    {
257        let _db = DB::open_default(&path).unwrap();
258        let opts = Options::default();
259        // The DB will still be open when we try to destroy it and the lock should fail.
260        match DB::destroy(&opts, &path) {
261            Err(s) => {
262                let message = s.to_string();
263                assert!(message.contains("IO error:"));
264                assert!(message.contains("/LOCK:"));
265            }
266            Ok(_) => panic!("should fail"),
267        }
268    }
269}
270
271#[test]
272fn writebatch_works() {
273    use crate::{prelude::*, TemporaryDBPath, WriteBatch};
274
275    let path = TemporaryDBPath::new();
276    {
277        let db = DB::open_default(&path).unwrap();
278        {
279            // test putx
280            let mut batch = WriteBatch::default();
281            assert!(db.get(b"k1").unwrap().is_none());
282            assert_eq!(batch.len(), 0);
283            assert!(batch.is_empty());
284            let _ = batch.put(b"k1", b"v1111");
285            let _ = batch.put(b"k2", b"v2222");
286            let _ = batch.put(b"k3", b"v3333");
287            assert_eq!(batch.len(), 3);
288            assert!(!batch.is_empty());
289            assert!(db.get(b"k1").unwrap().is_none());
290            let p = db.write(&batch);
291            assert!(p.is_ok());
292            let r: Result<Option<DBVector>, Error> = db.get(b"k1");
293            assert!(r.unwrap().unwrap().to_utf8().unwrap() == "v1111");
294        }
295        {
296            // test delete
297            let mut batch = WriteBatch::default();
298            let _ = batch.delete(b"k1");
299            assert_eq!(batch.len(), 1);
300            assert!(!batch.is_empty());
301            let p = db.write(&batch);
302            assert!(p.is_ok());
303            assert!(db.get(b"k1").unwrap().is_none());
304        }
305        {
306            // test delete_range
307            let mut batch = WriteBatch::default();
308            let _ = batch.delete_range(b"k2", b"k4");
309            assert_eq!(batch.len(), 1);
310            assert!(!batch.is_empty());
311            let p = db.write(&batch);
312            assert!(p.is_ok());
313            assert!(db.get(b"k2").unwrap().is_none());
314            assert!(db.get(b"k3").unwrap().is_none());
315        }
316        {
317            // test size_in_bytes
318            let mut batch = WriteBatch::default();
319            let before = batch.size_in_bytes();
320            let _ = batch.put(b"k1", b"v1234567890");
321            let after = batch.size_in_bytes();
322            assert!(before + 10 <= after);
323        }
324    }
325}
326
327#[test]
328fn iterator_test() {
329    use crate::{prelude::*, IteratorMode, TemporaryDBPath};
330    use std::str;
331
332    let path = TemporaryDBPath::new();
333    {
334        let db = DB::open_default(&path).unwrap();
335        let p = db.put(b"k1", b"v1111");
336        assert!(p.is_ok());
337        let p = db.put(b"k2", b"v2222");
338        assert!(p.is_ok());
339        let p = db.put(b"k3", b"v3333");
340        assert!(p.is_ok());
341        let iter = db.iterator(IteratorMode::Start);
342        for (k, v) in iter {
343            println!(
344                "Hello {}: {}",
345                str::from_utf8(&k).unwrap(),
346                str::from_utf8(&v).unwrap()
347            );
348        }
349    }
350}
351
352#[test]
353fn snapshot_test() {
354    use crate::{prelude::*, TemporaryDBPath};
355
356    let path = TemporaryDBPath::new();
357    {
358        let db = DB::open_default(&path).unwrap();
359        let p = db.put(b"k1", b"v1111");
360        assert!(p.is_ok());
361
362        let snap = db.snapshot();
363        let r: Result<Option<DBVector>, Error> = snap.get(b"k1");
364        assert!(r.unwrap().unwrap().to_utf8().unwrap() == "v1111");
365
366        let p = db.put(b"k2", b"v2222");
367        assert!(p.is_ok());
368
369        assert!(db.get(b"k2").unwrap().is_some());
370        assert!(snap.get(b"k2").unwrap().is_none());
371    }
372}
373
374#[test]
375fn set_option_test() {
376    use crate::{prelude::*, TemporaryDBPath};
377
378    let path = TemporaryDBPath::new();
379    {
380        let db = DB::open_default(&path).unwrap();
381        // set an option to valid values
382        assert!(db
383            .set_options(&[("disable_auto_compactions", "true")])
384            .is_ok());
385        assert!(db
386            .set_options(&[("disable_auto_compactions", "false")])
387            .is_ok());
388        // invalid names/values should result in an error
389        assert!(db
390            .set_options(&[("disable_auto_compactions", "INVALID_VALUE")])
391            .is_err());
392        assert!(db
393            .set_options(&[("INVALID_NAME", "INVALID_VALUE")])
394            .is_err());
395        // option names/values must not contain NULLs
396        assert!(db
397            .set_options(&[("disable_auto_compactions", "true\0")])
398            .is_err());
399        assert!(db
400            .set_options(&[("disable_auto_compactions\0", "true")])
401            .is_err());
402        // empty options are not allowed
403        assert!(db.set_options(&[]).is_err());
404        // multiple options can be set in a single API call
405        let multiple_options = [
406            ("paranoid_file_checks", "true"),
407            ("report_bg_io_stats", "true"),
408        ];
409        db.set_options(&multiple_options).unwrap();
410    }
411}