wip

crates/katana/storage/db/src/lib.rs

@@ -24,12 +24,12 @@ use version::{check_db_version, create_db_version_file, DatabaseVersionError};
 ///
 /// This will create the default tables, if necessary.
 pub fn init_db<P: AsRef<Path>>(path: P) -> anyhow::Result<DbEnv> {
-    init_db_with_schema::<tables::Tables>(path)
+    init_db_with_schema::<tables::SchemaV1>(path)
 }
 
 /// Open the database at the given `path` in read-write mode.
 pub fn open_db<P: AsRef<Path>>(path: P) -> anyhow::Result<DbEnv> {
-    open_db_with_schema::<tables::Tables>(path)
+    open_db_with_schema::<tables::SchemaV1>(path)
 }
 
 pub(crate) fn init_db_with_schema<S: Schema>(path: impl AsRef<Path>) -> anyhow::Result<DbEnv<S>> {

crates/katana/storage/db/src/mdbx/cursor.rs

@@ -3,7 +3,7 @@
 use std::borrow::Cow;
 use std::marker::PhantomData;
 
-use libmdbx::{self, TransactionKind, WriteFlags, RW};
+use libmdbx::{TransactionKind, WriteFlags, RW};
 
 use crate::codecs::{Compress, Encode};
 use crate::error::DatabaseError;

crates/katana/storage/db/src/mdbx/mod.rs

@@ -11,7 +11,7 @@ use libmdbx::{DatabaseFlags, EnvironmentFlags, Geometry, Mode, PageSize, SyncMod
 
 use self::tx::Tx;
 use crate::error::DatabaseError;
-use crate::tables::{Schema, TableType, Tables};
+use crate::tables::{Schema, SchemaV1, TableType};
 use crate::utils;
 
 const GIGABYTE: usize = 1024 * 1024 * 1024;
@@ -31,7 +31,10 @@ pub enum DbEnvKind {
 
 /// Wrapper for `libmdbx-sys` environment.
 #[derive(Debug)]
-pub struct DbEnv<S: Schema = Tables> {
+pub struct DbEnv<S = SchemaV1>
+where
+    S: Schema,
+{
     inner: libmdbx::Environment,
     _tables: std::marker::PhantomData<S>,
 }
@@ -93,20 +96,20 @@ impl<S: Schema> DbEnv<S> {
     }
 
     /// Begin a read-only transaction.
-    pub fn tx(&self) -> Result<Tx<RO>, DatabaseError> {
+    pub fn tx(&self) -> Result<Tx<RO, S>, DatabaseError> {
         Ok(Tx::new(self.inner.begin_ro_txn().map_err(DatabaseError::CreateROTx)?))
     }
 
     /// Begin a read-write transaction.
-    pub fn tx_mut(&self) -> Result<Tx<RW>, DatabaseError> {
+    pub fn tx_mut(&self) -> Result<Tx<RW, S>, DatabaseError> {
         Ok(Tx::new(self.inner.begin_rw_txn().map_err(DatabaseError::CreateRWTx)?))
     }
 
     /// Takes a function and passes a read-write transaction into it, making sure it's always
     /// committed in the end of the execution.
     pub fn update<T, F>(&self, f: F) -> Result<T, DatabaseError>
     where
-        F: FnOnce(&Tx<RW>) -> T,
+        F: FnOnce(&Tx<RW, S>) -> T,
     {
         let tx = self.tx_mut()?;
         let res = f(&tx);
@@ -118,6 +121,7 @@ impl<S: Schema> DbEnv<S> {
 #[cfg(any(test, feature = "test-utils"))]
 pub mod test_utils {
     use super::*;
+    use crate::tables;
 
     const ERROR_DB_CREATION: &str = "Not able to create the mdbx file.";
 
@@ -131,7 +135,7 @@ pub mod test_utils {
 
     /// Create database for testing with specified path
     pub fn create_test_db_with_path(kind: DbEnvKind, path: &Path) -> DbEnv {
-        let env = DbEnv::<Tables>::open(path, kind).expect(ERROR_DB_CREATION);
+        let env = DbEnv::<tables::SchemaV1>::open(path, kind).expect(ERROR_DB_CREATION);
         env.create_tables().expect("Failed to create tables.");
         env
     }

crates/katana/storage/db/src/mdbx/tx.rs

@@ -1,6 +1,6 @@
 //! Transaction wrapper for libmdbx-sys.
 
-use std::str::FromStr;
+use std::marker::PhantomData;
 
 use libmdbx::ffi::DBI;
 use libmdbx::{TransactionKind, WriteFlags, RW};
@@ -9,29 +9,38 @@ use parking_lot::RwLock;
 use super::cursor::Cursor;
 use crate::codecs::{Compress, Encode};
 use crate::error::DatabaseError;
-use crate::tables::{Table, Tables, NUM_TABLES};
+use crate::tables::{Schema, SchemaV1, Table, NUM_TABLES};
 use crate::utils::decode_one;
 
-/// Alias for read-only transaction.
-pub type TxRO = Tx<libmdbx::RO>;
-/// Alias for read-write transaction.
-pub type TxRW = Tx<libmdbx::RW>;
+/// Alias for read-only transaction on the default schema.
+pub type TxRO = Tx<libmdbx::RO, SchemaV1>;
+/// Alias for read-write transaction on the default schema.
+pub type TxRW = Tx<libmdbx::RW, SchemaV1>;
 
 /// Database transaction.
 ///
 /// Wrapper for a `libmdbx` transaction.
 #[derive(Debug)]
-pub struct Tx<K: TransactionKind> {
+pub struct Tx<K: TransactionKind, S: Schema> {
     /// Libmdbx-sys transaction.
     inner: libmdbx::Transaction<K>,
+    /// Marker for the db schema.
+    _schema: std::marker::PhantomData<S>,
+    // the array size is hardcoded to the number of tables in current db version for now. ideally
+    // we could use the associated constant from the schema trait. but that would require the
+    // `generic_const_exprs`.
     /// Database table handle cache.
     db_handles: RwLock<[Option<DBI>; NUM_TABLES]>,
 }
 
-impl<K: TransactionKind> Tx<K> {
+impl<K, S> Tx<K, S>
+where
+    K: TransactionKind,
+    S: Schema,
+{
     /// Creates new `Tx` object with a `RO` or `RW` transaction.
     pub fn new(inner: libmdbx::Transaction<K>) -> Self {
-        Self { inner, db_handles: Default::default() }
+        Self { inner, _schema: PhantomData, db_handles: Default::default() }
     }
 
     /// Creates a cursor to iterate over a table items.
@@ -44,13 +53,18 @@ impl<K: TransactionKind> Tx<K> {
 
     /// Gets a table database handle if it exists, otherwise creates it.
     pub fn get_dbi<T: Table>(&self) -> Result<DBI, DatabaseError> {
+        // SAFETY:
+        // the index is guaranteed to be in bounds by the schema only on current schema
+        // version because we hardcode the size exactly for the number of tables in current db
+        // schema. see `tables::v1::NUM_TABLES`.
+        let table = S::index::<T>().expect(&format!("table {} not found in schema", T::NAME));
+
         let mut handles = self.db_handles.write();
-        let table = Tables::from_str(T::NAME).expect("requested table should be part of `Tables`.");
+        let dbi_handle = handles.get_mut(table).expect("should exist");
 
-        let dbi_handle = handles.get_mut(table as usize).expect("should exist");
         if dbi_handle.is_none() {
-            *dbi_handle =
-                Some(self.inner.open_db(Some(T::NAME)).map_err(DatabaseError::OpenDb)?.dbi());
+            let dbi = self.inner.open_db(Some(T::NAME)).map_err(DatabaseError::OpenDb)?.dbi();
+            *dbi_handle = Some(dbi);
         }
 
         Ok(dbi_handle.expect("is some; qed"))
@@ -66,6 +80,22 @@ impl<K: TransactionKind> Tx<K> {
             .transpose()
     }
 
+    /// Gets a value from a table using the given key without checking if the table exist in the
+    /// schema.
+    pub fn get_unchecked<T: Table>(
+        &self,
+        key: T::Key,
+    ) -> Result<Option<<T as Table>::Value>, DatabaseError> {
+        let dbi = self.inner.open_db(Some(T::NAME)).map_err(DatabaseError::OpenDb)?.dbi();
+        let key = Encode::encode(key);
+
+        self.inner
+            .get(dbi, key.as_ref())
+            .map_err(DatabaseError::Read)?
+            .map(decode_one::<T>)
+            .transpose()
+    }
+
     /// Returns number of entries in the table using cheap DB stats invocation.
     pub fn entries<T: Table>(&self) -> Result<usize, DatabaseError> {
         self.inner
@@ -80,7 +110,7 @@ impl<K: TransactionKind> Tx<K> {
     }
 }
 
-impl Tx<RW> {
+impl<S: Schema> Tx<RW, S> {
     /// Inserts an item into a database.
     ///
     /// This function stores key/data pairs in the database. The default behavior is to enter the

crates/katana/storage/db/src/tables/mod.rs

@@ -1,10 +1,10 @@
 pub mod v0;
 mod v1;
 
-pub use self::v1::*;
 use std::fmt::{Debug, Display};
 use std::str::FromStr;
 
+pub use self::v1::*;
 use crate::codecs::{Compress, Decode, Decompress, Encode};
 
 pub trait Key: Encode + Decode + Clone + std::fmt::Debug {}
@@ -25,6 +25,8 @@ pub trait Schema: Debug + Display + FromStr + 'static {
     fn name(&self) -> &str;
     /// The type of the tables.
     fn table_type(&self) -> TableType;
+
+    fn index<T: Table>() -> Option<usize>;
 }
 
 /// An asbtraction for a table.
@@ -54,54 +56,63 @@ pub enum TableType {
     DupSort,
 }
 
-/// Macro to declare database tables based on the [DatabaseTables] trait.
+/// Macro to declare database tables based on the [Database$name] trait.
 #[macro_export]
-macro_rules! define_tables_enum {
-    { $ver:expr, [$(($table:ident, $type:expr)),*] } => {
+macro_rules! define_schema_enum {
+    { $ver:expr, $name:ident, [$(($table:ident, $type:expr)),*] } => {
         #[derive(Debug, PartialEq, Copy, Clone)]
-        pub enum Tables {
+        pub enum $name {
             $(
                 $table,
             )*
         }
 
-        impl Schema for Tables {
+        impl Schema for $name {
         	const VERSION: u32 = $ver;
 
-            fn all() -> &'static [Self] {
-				&[$(Tables::$table,)*]
-			}
+            fn all() -> &'static [Self] { &[$($name::$table,)*] }
 
             fn name(&self) -> &str {
                 match self {
-                    $(Tables::$table => {
+                    $($name::$table => {
                         $table::NAME
                     },)*
                 }
             }
 
             fn table_type(&self) -> TableType {
                 match self {
-                    $(Tables::$table => {
+                    $($name::$table => {
                         $type
                     },)*
                 }
             }
+
+            fn index<T: Table>() -> Option<usize> {
+	            match T::NAME {
+	                $($table::NAME => {
+	                    Some($name::$table as usize)
+	                },)*
+	                _ => {
+	                    None
+	                }
+	            }
+            }
         }
 
-        impl std::fmt::Display for Tables {
+        impl std::fmt::Display for $name {
             fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
                 write!(f, "{}", self.name())
             }
         }
 
-        impl std::str::FromStr for Tables {
+        impl std::str::FromStr for $name {
             type Err = String;
 
             fn from_str(s: &str) -> Result<Self, Self::Err> {
                 match s {
                     $($table::NAME => {
-                        return Ok(Tables::$table)
+                        return Ok($name::$table)
                     },)*
                     _ => {
                         return Err(format!("unknown table `{s}`"))

crates/katana/storage/db/src/tables/v0.rs

@@ -2,18 +2,18 @@ use katana_primitives::block::BlockNumber;
 use katana_primitives::contract::{ContractAddress, StorageKey};
 
 use super::*;
-use crate::codecs::{Compress, Decode, Decompress, Encode};
 use crate::error::CodecError;
 use crate::models::contract::{ContractClassChange, ContractNonceChange};
 use crate::models::storage::{ContractStorageEntry, ContractStorageKey};
-use crate::{define_tables_enum, dupsort, tables};
+use crate::{define_schema_enum, dupsort, tables};
 
 pub const NUM_TABLES: usize = 22;
 
 // TODO(kariy): can we somehow define this without repeating existing tables, and only add the older
 // ones?
-define_tables_enum! {
+define_schema_enum! {
     0,
+    SchemaV0,
     [(Headers, TableType::Table),
     (BlockHashes, TableType::Table),
     (BlockNumbers, TableType::Table),