1 files changed, 37 insertions, 3 deletions
diff --git a/src/state.rs b/src/state.rs
index 771647d..f7fe10e 100644
--- a/src/state.rs
+++ b/src/state.rs
@@ -1,17 +1,40 @@
-use std::fmt::Debug;
+use std::{fmt::Debug, path::Path, sync::Arc};
 
 use axum::extract::FromRef;
 use sqlx::PgPool;
 
 use crate::Error;
 
+pub type KVStore = rocksdb::OptimisticTransactionDB<rocksdb::MultiThreaded>;
+
 #[derive(Debug, Clone, FromRef)]
 pub struct AppState {
     pub pool: PgPool,
+    pub kv_store: Arc<rocksdb::OptimisticTransactionDB<rocksdb::MultiThreaded>>,
 }
 
 impl AppState {
-    pub async fn new(uri: String) -> Result<Self, Error> {
+    pub async fn new<P>(uri: String, path: P) -> Result<Self, Error>
+    where
+        P: AsRef<Path>,
+    {
+        Ok(Self {
+            pool: Self::init_pool(uri).await?,
+            kv_store: Arc::new(Self::init_kv_store(path).unwrap()),
+        })
+    }
+
+    pub fn with_pool<P>(pool: PgPool, path: P) -> Self
+    where
+        P: AsRef<Path>,
+    {
+        Self {
+            pool,
+            kv_store: Arc::new(Self::init_kv_store(path).unwrap()),
+        }
+    }
+
+    pub async fn init_pool(uri: String) -> Result<PgPool, sqlx::Error> {
         tracing::debug!("Attempting to connect to database...");
 
         let pool = sqlx::postgres::PgPoolOptions::new()
@@ -23,6 +46,17 @@ impl AppState {
 
         sqlx::migrate!().run(&pool).await?;
 
-        Ok(Self { pool })
+        Ok(pool)
+    }
+
+    pub fn init_kv_store<P>(
+        path: P,
+    ) -> Result<rocksdb::OptimisticTransactionDB<rocksdb::MultiThreaded>, rocksdb::Error>
+    where
+        P: AsRef<Path>,
+    {
+        let mut opts = rocksdb::Options::default();
+        opts.create_if_missing(true);
+        rocksdb::OptimisticTransactionDB::open(&opts, path)
     }
 }