feat: Added in-memory storage for testing purposes

Author: Harshdev098

rust/Cargo.lock

@@ -24,6 +24,8 @@ cargo build --release
    ```
    cargo run -- server/vss-server-config.toml
    ```
+
+   **Note:** For testing puropose you can edit vss-server-config.toml to use `store_type` as in-memory instead of postgresql `store_type = "memory"`
 4. VSS endpoint should be reachable at `http://localhost:8080/vss`.
 
 ### Configuration

rust/README.md

@@ -14,6 +14,8 @@
 mod migrations;
 /// Contains [PostgreSQL](https://www.postgresql.org/) based backend implementation for VSS.
 pub mod postgres_store;
+/// Contains in-memory backend implementation for VSS, for testing purposes only.
+pub mod memory_store;
 
 #[macro_use]
 extern crate api;

rust/impls/src/lib.rs

@@ -0,0 +1,317 @@
+use std::collections::HashMap;
+use std::sync::{Arc, Mutex};
+
+use chrono::prelude::{Utc};
+use async_trait::async_trait;
+use bytes::Bytes;
+
+use api::error::VssError;
+use crate::postgres_store::{VssDbRecord, LIST_KEY_VERSIONS_MAX_PAGE_SIZE, MAX_PUT_REQUEST_ITEM_COUNT};
+use api::kv_store::{KvStore, GLOBAL_VERSION_KEY};
+use api::types::{
+    DeleteObjectRequest, DeleteObjectResponse, GetObjectRequest, GetObjectResponse,
+    KeyValue, ListKeyVersionsRequest, ListKeyVersionsResponse, PutObjectRequest, PutObjectResponse,
+};
+
+/// In-memory backend for VSS, for testing purposes only.
+pub struct InMemoryBackendImpl {
+    store: Arc<Mutex<HashMap<(String, String, String), VssDbRecord>>>,
+}
+
+impl InMemoryBackendImpl {
+    /// Creates a new in-memory backend.
+    pub fn new() -> Self {
+        Self {
+            store: Arc::new(Mutex::new(HashMap::new())),
+        }
+    }
+
+    fn build_vss_record(&self, user_token: String, store_id: String, kv: KeyValue) -> VssDbRecord {
+        let now = Utc::now();
+        VssDbRecord {
+            user_token,
+            store_id,
+            key: kv.key,
+            value: kv.value.to_vec(),
+            version: kv.version,
+            created_at: now,
+            last_updated_at: now,
+        }
+    }
+
+    fn execute_put_object(
+        &self,
+        store: &mut HashMap<(String, String, String), VssDbRecord>,
+        record: &VssDbRecord,
+    ) -> Result<u64, VssError> {
+        let key_tuple = (
+            record.user_token.clone(),
+            record.store_id.clone(),
+            record.key.clone(),
+        );
+        let now = Utc::now();
+
+        if let Some(existing) = store.get_mut(&key_tuple) {
+            if existing.version >= record.version {
+                return Err(VssError::ConflictError(
+                    "Version conflict on put".to_string(),
+                ));
+            }
+            existing.version = record.version;
+            existing.value = record.value.clone();
+            existing.last_updated_at = now;
+            Ok(1)
+        } else {
+            store.insert(
+                key_tuple,
+                VssDbRecord {
+                    user_token: record.user_token.clone(),
+                    store_id: record.store_id.clone(),
+                    key: record.key.clone(),
+                    value: record.value.clone(),
+                    version: record.version,
+                    created_at: now,
+                    last_updated_at: now,
+                },
+            );
+            Ok(1)
+        }
+    }
+
+    fn execute_delete_object(
+        &self,
+        store: &mut HashMap<(String, String, String), VssDbRecord>,
+        record: &VssDbRecord,
+    ) -> Result<u64, VssError> {
+        let key_tuple = (
+            record.user_token.clone(),
+            record.store_id.clone(),
+            record.key.clone(),
+        );
+        if let Some(existing) = store.get(&key_tuple) {
+            if existing.version != record.version {
+                return Ok(0);
+            }
+            store.remove(&key_tuple);
+            Ok(1)
+        } else {
+            Ok(0)
+        }
+    }
+}
+
+#[async_trait]
+impl KvStore for InMemoryBackendImpl {
+    async fn get(
+        &self,
+        user_token: String,
+        request: GetObjectRequest,
+    ) -> Result<GetObjectResponse, VssError> {
+        let key_tuple = (user_token.clone(), request.store_id.clone(), request.key.clone());
+        let guard = self.store.lock().unwrap();
+
+        if let Some(record) = guard.get(&key_tuple) {
+            Ok(GetObjectResponse {
+                value: Some(KeyValue {
+                    key: record.key.clone(),
+                    value: Bytes::from(record.value.clone()),
+                    version: record.version,
+                }),
+            })
+        } else if request.key == GLOBAL_VERSION_KEY {
+            Ok(GetObjectResponse {
+                value: Some(KeyValue {
+                    key: GLOBAL_VERSION_KEY.to_string(),
+                    value: Bytes::new(),
+                    version: 0,
+                }),
+            })
+        } else {
+            Err(VssError::NoSuchKeyError("Requested key not found.".to_string()))
+        }
+    }
+
+    async fn put(
+        &self,
+        user_token: String,
+        request: PutObjectRequest,
+    ) -> Result<PutObjectResponse, VssError> {
+        if request.transaction_items.len() + request.delete_items.len() > MAX_PUT_REQUEST_ITEM_COUNT {
+            return Err(VssError::InvalidRequestError(format!(
+                "Number of write items per request should be less than equal to {}",
+                MAX_PUT_REQUEST_ITEM_COUNT
+            )));
+        }
+
+        let mut guard = self.store.lock().unwrap();
+        let mut temp_store = guard.clone();
+
+        let vss_put_records: Vec<VssDbRecord> = request
+            .transaction_items
+            .into_iter()
+            .map(|kv| self.build_vss_record(user_token.clone(), request.store_id.clone(), kv))
+            .collect();
+
+        let vss_delete_records: Vec<VssDbRecord> = request
+            .delete_items
+            .into_iter()
+            .map(|kv| self.build_vss_record(user_token.clone(), request.store_id.clone(), kv))
+            .collect();
+
+        let mut batch_results = Vec::new();
+
+        for vss_record in &vss_put_records {
+            let num_rows = self.execute_put_object(&mut temp_store, vss_record)?;
+            batch_results.push(num_rows);
+        }
+
+        for vss_record in &vss_delete_records {
+            let num_rows = self.execute_delete_object(&mut temp_store, vss_record)?;
+            batch_results.push(num_rows);
+        }
+
+        if batch_results.iter().any(|&rows| rows == 0 && !vss_delete_records.is_empty()) {
+            return Err(VssError::ConflictError(
+                "Transaction could not be completed due to a possible conflict".to_string(),
+            ));
+        }
+
+        *guard = temp_store;
+        Ok(PutObjectResponse {})
+    }
+
+    async fn delete(
+        &self,
+        user_token: String,
+        request: DeleteObjectRequest,
+    ) -> Result<DeleteObjectResponse, VssError> {
+        let key_value = request.key_value.ok_or_else(|| {
+            VssError::InvalidRequestError("key_value missing in DeleteObjectRequest".to_string())
+        })?;
+        let vss_record = self.build_vss_record(user_token, request.store_id, key_value);
+
+        let mut guard = self.store.lock().unwrap();
+        let mut temp_store = guard.clone();
+
+        self.execute_delete_object(&mut temp_store, &vss_record)?;
+
+        *guard = temp_store;
+        Ok(DeleteObjectResponse {})
+    }
+
+    async fn list_key_versions(
+        &self,
+        user_token: String,
+        request: ListKeyVersionsRequest,
+    ) -> Result<ListKeyVersionsResponse, VssError> {
+        let store_id = request.store_id;
+        let key_prefix = request.key_prefix.unwrap_or_default();
+        let page_token = request.page_token.unwrap_or_default();
+        let page_size = request.page_size.unwrap_or(i32::MAX);
+        let limit = std::cmp::min(page_size, LIST_KEY_VERSIONS_MAX_PAGE_SIZE) as usize;
+
+        let mut global_version = None;
+        if page_token.is_empty() {
+            let get_global_version_request = GetObjectRequest {
+                store_id: store_id.clone(),
+                key: GLOBAL_VERSION_KEY.to_string(),
+            };
+            let get_response = self.get(user_token.clone(), get_global_version_request).await?;
+            global_version = Some(get_response.value.unwrap().version);
+        }
+
+        let key_versions: Vec<KeyValue> = {
+            let guard = self.store.lock().unwrap();
+            let mut key_versions: Vec<KeyValue> = guard
+                .iter()
+                .filter(|(k, _)| {
+                    k.0 == user_token
+                        && k.1 == store_id
+                        && k.2.starts_with(&key_prefix)
+                        && k.2 > page_token
+                        && k.2 != GLOBAL_VERSION_KEY
+                })
+                .map(|(_, record)| KeyValue {
+                    key: record.key.clone(),
+                    value: Bytes::new(),
+                    version: record.version,
+                })
+                .collect();
+
+            key_versions.sort_by(|a, b| a.key.cmp(&b.key));
+            key_versions.into_iter().take(limit).collect()
+        };
+
+        let next_page_token = if key_versions.len() == limit {
+            key_versions.last().map(|kv| kv.key.clone())
+        } else {
+            None
+        };
+
+        Ok(ListKeyVersionsResponse {
+            key_versions,
+            next_page_token,
+            global_version,
+        })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use tokio::test;
+    use bytes::Bytes;
+    use api::kv_store::INITIAL_RECORD_VERSION;
+
+    #[test]
+    async fn test_in_memory_crud() {
+        let store = InMemoryBackendImpl::new();
+        let user_token = "test_user".to_string();
+        let store_id = "test_store".to_string();
+
+        // Put
+        let put_request = PutObjectRequest {
+            store_id: store_id.clone(),
+            transaction_items: vec![KeyValue {
+                key: "key1".to_string(),
+                value: Bytes::from("value1"),
+                version: INITIAL_RECORD_VERSION as i64,
+            }],
+            delete_items: vec![],
+            global_version: None,
+        };
+        store.put(user_token.clone(), put_request).await.unwrap();
+
+        // Get
+        let get_request = GetObjectRequest {
+            store_id: store_id.clone(),
+            key: "key1".to_string(),
+        };
+        let response = store.get(user_token.clone(), get_request).await.unwrap();
+        let key_value = response.value.unwrap();
+        assert_eq!(key_value.value, Bytes::from("value1"));
+        let current_version = key_value.version;
+
+        // List
+        let list_request = ListKeyVersionsRequest {
+            store_id: store_id.clone(),
+            key_prefix: None,
+            page_size: Some(1),
+            page_token: None,
+        };
+        let response = store.list_key_versions(user_token.clone(), list_request).await.unwrap();
+        assert_eq!(response.key_versions.len(), 1);
+        assert_eq!(response.key_versions[0].key, "key1");
+
+        // Delete
+        let delete_request = DeleteObjectRequest {
+            store_id,
+            key_value: Some(KeyValue {
+                key: "key1".to_string(),
+                value: Bytes::new(),
+                version: current_version,
+            }),
+        };
+        store.delete(user_token, delete_request).await.unwrap();
+    }
+}

rust/impls/src/memory_store.rs

@@ -5,7 +5,7 @@ pub(crate) const MIGRATION_LOG_COLUMN: &str = "upgrade_from";
 pub(crate) const CHECK_DB_STMT: &str = "SELECT 1 FROM pg_database WHERE datname = $1";
 pub(crate) const INIT_DB_CMD: &str = "CREATE DATABASE";
 #[cfg(test)]
-const DROP_DB_CMD: &str = "DROP DATABASE";
+pub(crate) const DROP_DB_CMD: &str = "DROP DATABASE";
 pub(crate) const GET_VERSION_STMT: &str = "SELECT db_version FROM vss_db_version;";
 pub(crate) const UPDATE_VERSION_STMT: &str = "UPDATE vss_db_version SET db_version=$1;";
 pub(crate) const LOG_MIGRATION_STMT: &str = "INSERT INTO vss_db_upgrades VALUES($1);";
@@ -36,4 +36,4 @@ pub(crate) const MIGRATIONS: &[&str] = &[
 	);",
 ];
 #[cfg(test)]
-const DUMMY_MIGRATION: &str = "SELECT 1 WHERE FALSE;";
+pub(crate) const DUMMY_MIGRATION: &str = "SELECT 1 WHERE FALSE;";

rust/impls/src/migrations.rs

@@ -16,6 +16,7 @@ use std::io;
 use std::io::{Error, ErrorKind};
 use tokio_postgres::{error, NoTls, Transaction};
 
+#[derive(Clone)]
 pub(crate) struct VssDbRecord {
 	pub(crate) user_token: String,
 	pub(crate) store_id: String,