fix: Setting RocksDB as default database, making db configurable, and removing redis from deps

crates/cli/src/cfg.rs

@@ -1,12 +1,18 @@
 use anyhow::{Context, Result};
-use clap::{Args, Parser, Subcommand};
+use clap::{Args, Parser, Subcommand, ValueEnum};
 use config::{builder::DefaultState, ConfigBuilder, File};
 use dirs::home_dir;
 use dotenvy::dotenv;
 use log::{error, warn};
 use prism_errors::{DataAvailabilityError, GeneralError};
 use prism_prover::webserver::WebServerConfig;
-use prism_storage::redis::RedisConfig;
+use prism_storage::{
+    database::StorageBackend,
+    inmemory::InMemoryDatabase,
+    redis::RedisConfig,
+    rocksdb::{RocksDBConfig, RocksDBConnection},
+    Database, RedisConnection,
+};
 use serde::{Deserialize, Serialize};
 use std::{fs, path::Path, sync::Arc};
 
@@ -30,8 +36,8 @@ pub struct CommandArgs {
     #[arg(short, long, default_value = "INFO")]
     log_level: String,
 
-    #[arg(short = 'r', long)]
-    redis_client: Option<String>,
+    #[arg(short = 'n', long, default_value = "local")]
+    network_name: Option<String>,
 
     #[arg(long)]
     verifying_key: Option<String>,
@@ -43,7 +49,10 @@ pub struct CommandArgs {
     verifying_key_algorithm: Option<String>,
 
     #[arg(long)]
-    config_path: Option<String>,
+    home_path: Option<String>,
+
+    #[command(flatten)]
+    database: DatabaseArgs,
 
     #[command(flatten)]
     celestia: CelestiaArgs,
@@ -74,7 +83,7 @@ struct CelestiaArgs {
     #[arg(long)]
     operation_namespace_id: Option<String>,
 
-    // Height to start searching the DA layer for SNARKs on
+    /// Height to start searching the DA layer for SNARKs on
     #[arg(short = 's', long)]
     celestia_start_height: Option<u64>,
 }
@@ -101,18 +110,18 @@ pub struct Config {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub celestia_config: Option<CelestiaConfig>,
     pub da_layer: DALayerOption,
-    pub redis_config: Option<RedisConfig>,
+    pub db: StorageBackend,
     pub verifying_key: Option<String>,
     pub verifying_key_algorithm: String,
 }
 
-impl Default for Config {
-    fn default() -> Self {
+impl Config {
+    fn with_home(path: &str) -> Self {
         Config {
             webserver: Some(WebServerConfig::default()),
             celestia_config: Some(CelestiaConfig::default()),
             da_layer: DALayerOption::default(),
-            redis_config: Some(RedisConfig::default()),
+            db: StorageBackend::RocksDB(RocksDBConfig::new(path)),
             verifying_key: None,
             verifying_key_algorithm: "ed25519".to_string(),
         }
@@ -126,16 +135,38 @@ pub enum DALayerOption {
     InMemory,
 }
 
+#[derive(Debug, Default, Clone, Eq, PartialEq, Serialize, Deserialize, ValueEnum)]
+pub enum DBValues {
+    #[default]
+    RocksDB,
+    InMemory,
+    Redis,
+}
+
+#[derive(Args, Deserialize, Clone, Debug)]
+pub struct DatabaseArgs {
+    #[arg(long, value_enum)]
+    db_type: Option<DBValues>,
+
+    /// Path to the RocksDB database, used when `db_type` is `rocks-db`
+    #[arg(long)]
+    rocksdb_path: Option<String>,
+
+    /// Connection string to Redis, used when `db_type` is `redis`
+    #[arg(long)]
+    redis_url: Option<String>,
+}
+
 pub fn load_config(args: CommandArgs) -> Result<Config> {
     dotenv().ok();
     std::env::set_var("RUST_LOG", args.clone().log_level);
     pretty_env_logger::init();
 
-    let config_path = get_config_path(&args).context("Failed to determine config path")?;
-    ensure_config_file_exists(&config_path).context("Failed to ensure config file exists")?;
+    let home_path = get_prism_home(&args).context("Failed to determine prism home path")?;
+    ensure_config_file_exists(&home_path).context("Failed to ensure config file exists")?;
 
     let config_source = ConfigBuilder::<DefaultState>::default()
-        .add_source(File::with_name(&config_path))
+        .add_source(File::with_name(&format!("{}/config.toml", home_path)))
         .build()
         .context("Failed to build config")?;
 
@@ -151,22 +182,26 @@ pub fn load_config(args: CommandArgs) -> Result<Config> {
     Ok(final_config)
 }
 
-fn get_config_path(args: &CommandArgs) -> Result<String> {
-    args.config_path
+fn get_prism_home(args: &CommandArgs) -> Result<String> {
+    let network_name = args.network_name.clone().unwrap_or_else(|| "custom".to_string());
+    args.home_path
         .clone()
-        .or_else(|| home_dir().map(|path| format!("{}/.prism/config.toml", path.to_string_lossy())))
+        .or_else(|| {
+            home_dir().map(|path| format!("{}/.prism/{}/", path.to_string_lossy(), network_name))
+        })
         .ok_or_else(|| {
             GeneralError::MissingArgumentError("could not determine config path".to_string()).into()
         })
 }
 
-fn ensure_config_file_exists(config_path: &str) -> Result<()> {
+fn ensure_config_file_exists(home_path: &str) -> Result<()> {
+    let config_path = &format!("{}/config.toml", home_path);
     if !Path::new(config_path).exists() {
         if let Some(parent) = Path::new(config_path).parent() {
             fs::create_dir_all(parent).context("Failed to create config directory")?;
         }
 
-        let default_config = Config::default();
+        let default_config = Config::with_home(home_path);
         let config_toml =
             toml::to_string(&default_config).context("Failed to serialize default config")?;
 
@@ -177,18 +212,24 @@ fn ensure_config_file_exists(config_path: &str) -> Result<()> {
 
 fn apply_command_line_args(config: Config, args: CommandArgs) -> Config {
     let webserver_config = &config.webserver.unwrap_or_default();
-    let redis_config = &config.redis_config.unwrap_or_default();
     let celestia_config = &config.celestia_config.unwrap_or_default();
+    let prism_home = get_prism_home(&args.clone()).unwrap();
 
     Config {
         webserver: Some(WebServerConfig {
             enabled: args.webserver.webserver_active.unwrap_or(webserver_config.enabled),
             host: args.webserver.host.unwrap_or(webserver_config.host.clone()),
             port: args.webserver.port.unwrap_or(webserver_config.port),
         }),
-        redis_config: Some(RedisConfig {
-            connection_string: args.redis_client.unwrap_or(redis_config.connection_string.clone()),
-        }),
+        db: match args.database.db_type {
+            None | Some(DBValues::RocksDB) => StorageBackend::RocksDB(RocksDBConfig {
+                path: args.database.rocksdb_path.unwrap_or(prism_home),
+            }),
+            Some(DBValues::Redis) => StorageBackend::Redis(RedisConfig {
+                connection_string: args.database.redis_url.unwrap_or_default(),
+            }),
+            Some(DBValues::InMemory) => StorageBackend::InMemory,
+        },
         celestia_config: Some(CelestiaConfig {
             connection_string: args
                 .celestia
@@ -215,6 +256,27 @@ fn apply_command_line_args(config: Config, args: CommandArgs) -> Config {
     }
 }
 
+pub fn initialize_db(cfg: &Config) -> Result<Arc<Box<dyn Database>>> {
+    match &cfg.db {
+        StorageBackend::RocksDB(cfg) => {
+            let db = RocksDBConnection::new(cfg)
+                .map_err(|e| GeneralError::InitializationError(e.to_string()))
+                .context("Failed to initialize RocksDB")?;
+
+            Ok(Arc::new(Box::new(db) as Box<dyn Database>))
+        }
+        StorageBackend::InMemory => Ok(Arc::new(
+            Box::new(InMemoryDatabase::new()) as Box<dyn Database>
+        )),
+        StorageBackend::Redis(cfg) => {
+            let db = RedisConnection::new(cfg)
+                .map_err(|e| GeneralError::InitializationError(e.to_string()))
+                .context("Failed to initialize Redis")?;
+            Ok(Arc::new(Box::new(db) as Box<dyn Database>))
+        }
+    }
+}
+
 pub async fn initialize_da_layer(
     config: &Config,
 ) -> Result<Arc<dyn DataAvailabilityLayer + 'static>> {

crates/cli/src/main.rs

@@ -1,7 +1,7 @@
 mod cfg;
 mod node_types;
 
-use cfg::{initialize_da_layer, load_config, Cli, Commands};
+use cfg::{initialize_da_layer, initialize_db, load_config, Cli, Commands};
 use clap::Parser;
 use keystore_rs::{KeyChain, KeyStore, KeyStoreType};
 use prism_keys::{CryptoAlgorithm, SigningKey, VerifyingKey};
@@ -11,7 +11,6 @@ use std::io::{Error, ErrorKind};
 use node_types::NodeType;
 use prism_lightclient::LightClient;
 use prism_prover::Prover;
-use prism_storage::RedisConnection;
 use std::{str::FromStr, sync::Arc};
 
 #[macro_use]
@@ -67,12 +66,8 @@ async fn main() -> std::io::Result<()> {
                 .await
                 .map_err(|e| Error::new(ErrorKind::Other, e.to_string()))?;
 
-            let redis_config = config
-                .clone()
-                .redis_config
-                .ok_or_else(|| Error::new(ErrorKind::NotFound, "redis configuration not found"))?;
-            let redis_connections = RedisConnection::new(&redis_config)
-                .map_err(|e| Error::new(ErrorKind::Other, e.to_string()))?;
+            let db =
+                initialize_db(&config).map_err(|e| Error::new(ErrorKind::Other, e.to_string()))?;
 
             let signing_key_chain = KeyStoreType::KeyChain(KeyChain)
                 .get_signing_key()
@@ -106,14 +101,10 @@ async fn main() -> std::io::Result<()> {
                 start_height: config.celestia_config.unwrap_or_default().start_height,
             };
 
-            Arc::new(
-                Prover::new(Arc::new(Box::new(redis_connections)), da, &prover_cfg).map_err(
-                    |e| {
-                        error!("error initializing prover: {}", e);
-                        Error::new(ErrorKind::Other, e.to_string())
-                    },
-                )?,
-            )
+            Arc::new(Prover::new(db, da, &prover_cfg).map_err(|e| {
+                error!("error initializing prover: {}", e);
+                Error::new(ErrorKind::Other, e.to_string())
+            })?)
         }
         Commands::FullNode(args) => {
             let config = load_config(args.clone())
@@ -123,12 +114,8 @@ async fn main() -> std::io::Result<()> {
                 .await
                 .map_err(|e| Error::new(ErrorKind::Other, e.to_string()))?;
 
-            let redis_config = config
-                .clone()
-                .redis_config
-                .ok_or_else(|| Error::new(ErrorKind::NotFound, "redis configuration not found"))?;
-            let redis_connections = RedisConnection::new(&redis_config)
-                .map_err(|e| Error::new(ErrorKind::Other, e.to_string()))?;
+            let db =
+                initialize_db(&config).map_err(|e| Error::new(ErrorKind::Other, e.to_string()))?;
 
             let signing_key_chain = KeyStoreType::KeyChain(KeyChain)
                 .get_signing_key()
@@ -169,14 +156,10 @@ async fn main() -> std::io::Result<()> {
                 start_height: config.celestia_config.unwrap_or_default().start_height,
             };
 
-            Arc::new(
-                Prover::new(Arc::new(Box::new(redis_connections)), da, &prover_cfg).map_err(
-                    |e| {
-                        error!("error initializing prover: {}", e);
-                        Error::new(ErrorKind::Other, e.to_string())
-                    },
-                )?,
-            )
+            Arc::new(Prover::new(db, da, &prover_cfg).map_err(|e| {
+                error!("error initializing prover: {}", e);
+                Error::new(ErrorKind::Other, e.to_string())
+            })?)
         }
     };
 

crates/storage/src/database.rs

@@ -3,6 +3,14 @@ use auto_impl::auto_impl;
 use jmt::storage::{TreeReader, TreeWriter};
 use prism_common::digest::Digest;
 use prism_errors::{DatabaseError, PrismError};
+use serde::{Deserialize, Serialize};
+
+#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
+pub enum StorageBackend {
+    RocksDB(crate::rocksdb::RocksDBConfig),
+    InMemory,
+    Redis(crate::redis::RedisConfig),
+}
 
 #[auto_impl(&, Box, Arc)]
 pub trait Database: Send + Sync + TreeReader + TreeWriter {

crates/storage/src/redis.rs

@@ -23,7 +23,7 @@ use prism_errors::DatabaseError;
 use crate::database::{convert_to_connection_error, Database};
 use log::debug;
 
-#[derive(Debug, Serialize, Deserialize, Clone)]
+#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)]
 pub struct RedisConfig {
     pub connection_string: String,
 }

crates/storage/src/rocksdb.rs

@@ -13,21 +13,36 @@ use prism_serde::{
     hex::{FromHex, ToHex},
 };
 use rocksdb::{DBWithThreadMode, MultiThreaded, Options, DB};
+use serde::{Deserialize, Serialize};
 
 const KEY_PREFIX_COMMITMENTS: &str = "commitments:epoch_";
 const KEY_PREFIX_NODE: &str = "node:";
 const KEY_PREFIX_VALUE_HISTORY: &str = "value_history:";
 
 type RocksDB = DBWithThreadMode<MultiThreaded>;
 
+#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)]
+pub struct RocksDBConfig {
+    pub path: String,
+}
+
+impl RocksDBConfig {
+    pub fn new(path: &str) -> Self {
+        Self {
+            path: path.to_string(),
+        }
+    }
+}
+
 #[derive(Clone)]
 pub struct RocksDBConnection {
     connection: Arc<RocksDB>,
     path: String,
 }
 
 impl RocksDBConnection {
-    pub fn new(path: &str) -> Result<RocksDBConnection> {
+    pub fn new(cfg: &RocksDBConfig) -> Result<RocksDBConnection> {
+        let path = &cfg.path;
         let db = DB::open_default(path)?;
 
         Ok(Self {
@@ -180,7 +195,8 @@ mod tests {
 
     fn setup_db() -> (TempDir, RocksDBConnection) {
         let temp_dir = TempDir::new().unwrap();
-        let db = RocksDBConnection::new(temp_dir.path().to_str().unwrap()).unwrap();
+        let cfg = RocksDBConfig::new(temp_dir.path().to_str().unwrap());
+        let db = RocksDBConnection::new(&cfg).unwrap();
         (temp_dir, db)
     }
 

crates/tests/src/lib.rs

@@ -12,7 +12,10 @@ use prism_da::{
 use prism_keys::{CryptoAlgorithm, SigningKey};
 use prism_lightclient::LightClient;
 use prism_prover::Prover;
-use prism_storage::{rocksdb::RocksDBConnection, Database};
+use prism_storage::{
+    rocksdb::{RocksDBConfig, RocksDBConnection},
+    Database,
+};
 use rand::{rngs::StdRng, Rng, SeedableRng};
 use sp1_sdk::{HashableKey, ProverClient};
 use std::sync::Arc;
@@ -24,7 +27,8 @@ pub const PRISM_ELF: &[u8] = include_bytes!("../../../elf/riscv32im-succinct-zkv
 
 fn setup_db() -> Arc<Box<dyn Database>> {
     let temp_dir = TempDir::new().unwrap();
-    let db = RocksDBConnection::new(temp_dir.path().to_str().unwrap()).unwrap();
+    let cfg = RocksDBConfig::new(temp_dir.path().to_str().unwrap());
+    let db = RocksDBConnection::new(&cfg).unwrap();
     Arc::new(Box::new(db) as Box<dyn Database>)
 }