sov_db/
native_db.rs

1use std::path::Path;
2use std::sync::Arc;
3
4use sov_schema_db::{SchemaBatch, DB};
5
6use crate::rocks_db_config::gen_rocksdb_options;
7use crate::schema::tables::{ModuleAccessoryState, NATIVE_TABLES};
8use crate::schema::types::StateKey;
9
10/// A typed wrapper around RocksDB for storing native-only accessory state.
11/// Internally, this is roughly just an [`Arc<SchemaDB>`].
12#[derive(Clone, Debug)]
13pub struct NativeDB {
14    /// The underlying RocksDB instance, wrapped in an [`Arc`] for convenience
15    /// and [`DB`] for type safety.
16    db: Arc<DB>,
17}
18
19impl NativeDB {
20    const DB_PATH_SUFFIX: &'static str = "native";
21    const DB_NAME: &'static str = "native-db";
22
23    /// Opens a [`NativeDB`] (backed by RocksDB) at the specified path.
24    /// The returned instance will be at the path `{path}/native-db`.
25    pub fn with_path(path: impl AsRef<Path>) -> Result<Self, anyhow::Error> {
26        let path = path.as_ref().join(Self::DB_PATH_SUFFIX);
27        let inner = DB::open(
28            path,
29            Self::DB_NAME,
30            NATIVE_TABLES.iter().copied(),
31            &gen_rocksdb_options(&Default::default(), false),
32        )?;
33
34        Ok(Self {
35            db: Arc::new(inner),
36        })
37    }
38
39    /// Queries for a value in the [`NativeDB`], given a key.
40    pub fn get_value_option(&self, key: &StateKey) -> anyhow::Result<Option<Vec<u8>>> {
41        self.db
42            .get::<ModuleAccessoryState>(key)
43            .map(Option::flatten)
44    }
45
46    /// Sets a sequence of key-value pairs in the [`NativeDB`]. The write is atomic.
47    pub fn set_values(
48        &self,
49        key_value_pairs: impl IntoIterator<Item = (Vec<u8>, Option<Vec<u8>>)>,
50    ) -> anyhow::Result<()> {
51        let batch = SchemaBatch::default();
52        for (key, value) in key_value_pairs {
53            batch.put::<ModuleAccessoryState>(&key, &value)?;
54        }
55        self.db.write_schemas(batch)
56    }
57}
58
59#[cfg(feature = "arbitrary")]
60pub mod arbitrary {
61    //! Arbitrary definitions for the [`NativeDB`].
62
63    use core::ops::{Deref, DerefMut};
64
65    use proptest::strategy::LazyJust;
66    use tempfile::TempDir;
67
68    use super::*;
69
70    /// Arbitrary instance of [`NativeDB`].
71    ///
72    /// This is a db wrapper bound to its temporary path that will be deleted once the type is
73    /// dropped.
74    #[derive(Debug)]
75    pub struct ArbitraryNativeDB {
76        /// The underlying RocksDB instance.
77        pub db: NativeDB,
78        /// The temporary directory used to create the [`NativeDB`].
79        pub path: TempDir,
80    }
81
82    /// A fallible, arbitrary instance of [`NativeDB`].
83    ///
84    /// This type is suitable for operations that are expected to be infallible. The internal
85    /// implementation of the db requires I/O to create the temporary dir, making it fallible.
86    #[derive(Debug)]
87    pub struct FallibleArbitraryNativeDB {
88        /// The result of the new db instance.
89        pub result: anyhow::Result<ArbitraryNativeDB>,
90    }
91
92    impl Deref for ArbitraryNativeDB {
93        type Target = NativeDB;
94
95        fn deref(&self) -> &Self::Target {
96            &self.db
97        }
98    }
99
100    impl DerefMut for ArbitraryNativeDB {
101        fn deref_mut(&mut self) -> &mut Self::Target {
102            &mut self.db
103        }
104    }
105
106    impl<'a> ::arbitrary::Arbitrary<'a> for ArbitraryNativeDB {
107        fn arbitrary(_u: &mut ::arbitrary::Unstructured<'a>) -> ::arbitrary::Result<Self> {
108            let path = TempDir::new().map_err(|_| ::arbitrary::Error::NotEnoughData)?;
109            let db = NativeDB::with_path(&path).map_err(|_| ::arbitrary::Error::IncorrectFormat)?;
110            Ok(Self { db, path })
111        }
112    }
113
114    impl proptest::arbitrary::Arbitrary for FallibleArbitraryNativeDB {
115        type Parameters = ();
116        type Strategy = LazyJust<Self, fn() -> FallibleArbitraryNativeDB>;
117
118        fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
119            fn gen() -> FallibleArbitraryNativeDB {
120                FallibleArbitraryNativeDB {
121                    result: TempDir::new()
122                        .map_err(|e| {
123                            anyhow::anyhow!(format!("failed to generate path for NativeDB: {e}"))
124                        })
125                        .and_then(|path| {
126                            let db = NativeDB::with_path(&path)?;
127                            Ok(ArbitraryNativeDB { db, path })
128                        }),
129                }
130            }
131            LazyJust::new(gen)
132        }
133    }
134}
135
136#[cfg(test)]
137mod tests {
138    use super::*;
139
140    #[test]
141    fn get_after_set() {
142        let tmpdir = tempfile::tempdir().unwrap();
143        let db = NativeDB::with_path(tmpdir.path()).unwrap();
144
145        let key = b"foo".to_vec();
146        let value = b"bar".to_vec();
147        db.set_values(vec![(key.clone(), Some(value.clone()))])
148            .unwrap();
149        assert_eq!(db.get_value_option(&key).unwrap(), Some(value));
150    }
151
152    #[test]
153    fn get_after_delete() {
154        let tmpdir = tempfile::tempdir().unwrap();
155        let db = NativeDB::with_path(tmpdir.path()).unwrap();
156
157        let key = b"deleted".to_vec();
158        db.set_values(vec![(key.clone(), None)]).unwrap();
159        assert_eq!(db.get_value_option(&key).unwrap(), None);
160    }
161
162    #[test]
163    fn get_nonexistent() {
164        let tmpdir = tempfile::tempdir().unwrap();
165        let db = NativeDB::with_path(tmpdir.path()).unwrap();
166
167        let key = b"spam".to_vec();
168        assert_eq!(db.get_value_option(&key).unwrap(), None);
169    }
170}
sov_db/native_db.rs

sov_db/
native_db.rs
