Implement remaining rpc tests; Minor bug fixes

Author: TanmayArya-1p

Cargo.lock

@@ -8,3 +8,11 @@ pub enum DbError {
     LockError,
     DimensionMismatch,
 }
+
+impl std::fmt::Display for DbError {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        write!(f, "{:?}", self)
+    }
+}
+
+impl std::error::Error for DbError {}

crates/defs/src/error.rs

@@ -13,15 +13,16 @@ cargo run --bin grpc_server
 Use the [.sample.env](.sample.env) shown below as a reference to set your environment variables in a `.env` file.
 
 ```bash
-GRPC_SERVER_ROOT_PASSWORD=123 // required
-GRPC_SERVER_DIMENSION=3 // required
-
-GRPC_SERVER_HOST=localhost // defaults to 127.0.0.1 aka localhost
-GRPC_SERVER_PORT=8080 // defaults to 8080
-GRPC_SERVER_STORAGE_TYPE=inmemory // (inmemory/rocksdb) defaults to inmemory
-GRPC_SERVER_INDEX_TYPE=flat // defaults to flat
-GRPC_SERVER_DATA_PATH=data // defaults to a temporary directory
-GRPC_SERVER_LOGGING=true // defaults to true
+GRPC_SERVER_ROOT_PASSWORD=123 # required
+GRPC_SERVER_DIMENSION=3 # required
+
+GRPC_SERVER_HOST=localhost # defaults to 127.0.0.1 aka localhost
+GRPC_SERVER_PORT=8080 # defaults to 8080
+GRPC_SERVER_STORAGE_TYPE=inmemory # (inmemory/rocksdb) defaults to 'inmemory'
+GRPC_SERVER_INDEX_TYPE=flat # defaults to flat
+GRPC_SERVER_DATA_PATH=data # defaults to a temporary directory
+GRPC_SERVER_LOGGING=true # defaults to true
+
 
 ```
 

crates/grpc_server/README.md

@@ -110,7 +110,7 @@ impl GRPCServerConfig {
         let addr: SocketAddr = format!("{}:{}", host, port).parse()?;
 
         // check if logging is enabled
-        let mut logging: bool = true; // default to no logging
+        let mut logging: bool = true; // default to logging enabled
         if let Ok(logging_str) = env::var(ENV_LOGGING) {
             logging = logging_str.parse().unwrap_or(true);
         }

crates/grpc_server/src/config.rs

@@ -5,6 +5,8 @@ use std::panic;
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
     tracing_subscriber::fmt::init();
+
+    // load config from environment from environment variables
     let config = GRPCServerConfig::load_config()
         .inspect_err(|err| panic!("Failed to load config: {}", err))
         .unwrap();

crates/grpc_server/src/main.rs

@@ -1,7 +1,7 @@
 use crate::constants::SIMILARITY_PROTOBUFF_MAP;
 use crate::interceptors;
 use crate::utils::log_rpc;
-use std::{net::SocketAddr, panic};
+use std::net::SocketAddr;
 use tonic::{Request, Response, Status, service::InterceptorLayer, transport::Server};
 use tracing::{Level, event};
 use vectordb::{
@@ -82,12 +82,14 @@ impl VectorDb for VectorDBService {
         let query_vect = search_request
             .query_vector
             .ok_or(Status::invalid_argument("Invalid query_vector"))?;
-        let similarity = SIMILARITY_PROTOBUFF_MAP[search_request.similarity as usize];
+        let similarity = SIMILARITY_PROTOBUFF_MAP
+            .get(search_request.similarity as usize)
+            .ok_or(Status::internal("Invalid similarity"))?;
         let limit = search_request.limit;
 
         let result_point_ids = self
             .vector_db
-            .search(query_vect.values, similarity, limit as usize)
+            .search(query_vect.values, *similarity, limit as usize)
             .map_err(|_| Status::internal("Internal server error"))?;
 
         // create a mapped vector of PointIds
@@ -114,7 +116,7 @@ impl VectorDb for VectorDBService {
                     Err(Status::not_found("Point not found"))
                 }
             }
-            Err(_) => Err(Status::not_found("Error deleting point")),
+            Err(_) => Err(Status::internal("Error deleting point")),
         }
     }
 }
@@ -128,11 +130,14 @@ pub async fn run_server(
 
     let auth_interceptor = interceptors::AuthInterceptor::new(root_password);
 
-    let _ = Server::builder()
+    Server::builder()
         .layer(InterceptorLayer::new(auth_interceptor))
         .add_service(VectorDbServer::new(vector_db_service))
         .serve(addr)
         .await
-        .inspect_err(|err| panic!("Failed to start gRPC server: {}", err));
+        .map_err(|e| {
+            event!(Level::ERROR, "Failed to start gRPC server: {}", e);
+            Box::new(e)
+        })?;
     Ok(())
 }

crates/grpc_server/src/service.rs

@@ -3,7 +3,7 @@ use crate::interceptors;
 use crate::service::VectorDBService;
 use crate::service::vectordb::vector_db_client::VectorDbClient;
 use crate::service::vectordb::vector_db_server::VectorDbServer;
-use crate::service::vectordb::{DenseVector, InsertVectorRequest, PointId};
+use crate::service::vectordb::{DenseVector, InsertVectorRequest, PointId, SearchRequest};
 use api;
 use api::DbConfig;
 use index::IndexType;
@@ -22,7 +22,7 @@ use tonic::transport::{Channel, Server};
 // - use a shared instance of the server
 // currently tests must be run with --test-threads=1
 
-pub async fn run_test_server_with_stream(
+pub async fn run_server_with_stream(
     vector_db_service: VectorDBService,
     listener: TcpListener,
     root_password: String,
@@ -39,6 +39,7 @@ pub async fn run_test_server_with_stream(
 }
 
 async fn start_test_server() -> Result<(), Box<dyn std::error::Error>> {
+    // using a temporary directory for db datapath
     let temp_dir = tempdir().unwrap();
 
     let db_config = DbConfig {
@@ -65,8 +66,7 @@ async fn start_test_server() -> Result<(), Box<dyn std::error::Error>> {
     let listener = tokio::net::TcpListener::bind(config.addr).await?;
 
     tokio::spawn(async move {
-        let _ =
-            run_test_server_with_stream(vector_db_service, listener, config.root_password).await;
+        let _ = run_server_with_stream(vector_db_service, listener, config.root_password).await;
     });
 
     Ok(())
@@ -82,12 +82,14 @@ async fn create_test_client() -> Result<VectorDbClient<Channel>, Box<dyn std::er
 #[tokio::test]
 async fn test_grpc_server_start() {
     start_test_server().await.unwrap();
-
     let mut client = create_test_client().await.unwrap();
 
+    // insert a test vector
+    let test_vec = vec![1.0, 2.0, 3.0];
+
     let mut request = tonic::Request::new(InsertVectorRequest {
         vector: Some(DenseVector {
-            values: vec![1.0, 2.0, 3.0],
+            values: test_vec.clone(),
         }),
         payload: Some(Struct::default()),
     });
@@ -103,6 +105,7 @@ async fn test_insert_vector_rpc() {
     start_test_server().await.unwrap();
     let mut client = create_test_client().await.unwrap();
 
+    // insert a test vector
     let test_vec = vec![1.0, 2.0, 3.0];
 
     let mut request = tonic::Request::new(InsertVectorRequest {
@@ -116,10 +119,10 @@ async fn test_insert_vector_rpc() {
         .insert("authorization", "Bearer 123".parse().unwrap());
     let resp = client.insert_vector(request).await;
 
-    // Check if request is successful
+    // check if request is successful
     assert!(resp.is_ok());
 
-    // Check if the vector is actually present in the database
+    // check if the vector is actually present in the database
     let mut request = tonic::Request::new(PointId {
         id: resp.unwrap().into_inner().id,
     });
@@ -128,12 +131,12 @@ async fn test_insert_vector_rpc() {
         .insert("authorization", "Bearer 123".parse().unwrap());
     let resp = client.get_point(request).await;
 
-    // Check if request is successful
+    // check if request is successful
     assert!(resp.is_ok());
     let point = resp.unwrap().into_inner();
     assert_eq!(point.vector.unwrap().values, test_vec);
 
-    // Insert a new vector with mismatched dimensions
+    // insert a new vector with mismatched dimensions
     let mut request = tonic::Request::new(InsertVectorRequest {
         vector: Some(DenseVector {
             values: vec![1.0, 2.0],
@@ -145,15 +148,119 @@ async fn test_insert_vector_rpc() {
         .insert("authorization", "Bearer 123".parse().unwrap());
     let resp = client.insert_vector(request).await;
 
-    // Request must have failed
+    // request must fail
     assert!(resp.is_err());
 }
 
 #[tokio::test]
-async fn test_delete_vector_rpc() {}
+async fn test_delete_vector_rpc() {
+    start_test_server().await.unwrap();
+    let mut client = create_test_client().await.unwrap();
+
+    // insert a test vector
+    let test_vec = vec![1.0, 2.0, 3.0];
+    let mut request = tonic::Request::new(InsertVectorRequest {
+        vector: Some(DenseVector {
+            values: test_vec.clone(),
+        }),
+        payload: Some(Struct::default()),
+    });
+    request
+        .metadata_mut()
+        .insert("authorization", "Bearer 123".parse().unwrap());
+    let resp = client.insert_vector(request).await;
+
+    // check if request is successful
+    assert!(resp.is_ok());
+    let point = resp.unwrap().into_inner();
+
+    // delete the vector
+    let mut request = tonic::Request::new(PointId { id: point.id });
+    request
+        .metadata_mut()
+        .insert("authorization", "Bearer 123".parse().unwrap());
+    let resp = client.delete_point(request).await;
+
+    // check if request is successful
+    assert!(resp.is_ok());
+
+    // verify that the vector is deleted
+    let mut request = tonic::Request::new(PointId { id: point.id });
+    request
+        .metadata_mut()
+        .insert("authorization", "Bearer 123".parse().unwrap());
+    let resp = client.get_point(request).await;
+
+    // request must fail since the vector is deleted
+    assert!(resp.is_err());
+}
 
 #[tokio::test]
-async fn test_search_vector_rpc() {}
+async fn test_search_vector_rpc() {
+    start_test_server().await.unwrap();
+    let mut client = create_test_client().await.unwrap();
+
+    // insert a test vector
+    let test_vec = vec![1.0, 2.0, 3.0];
+    let mut request = tonic::Request::new(InsertVectorRequest {
+        vector: Some(DenseVector {
+            values: test_vec.clone(),
+        }),
+        payload: Some(Struct::default()),
+    });
+    request
+        .metadata_mut()
+        .insert("authorization", "Bearer 123".parse().unwrap());
+    let resp = client.insert_vector(request).await;
+
+    // check if request is successful
+    assert!(resp.is_ok());
+    let point = resp.unwrap().into_inner();
+
+    let query_vec = vec![2.0, 2.0, 2.0];
+
+    // search for the vector
+    let mut request = tonic::Request::new(SearchRequest {
+        query_vector: Some(DenseVector {
+            values: query_vec.clone(),
+        }),
+        similarity: 0, // euclidean distance
+        limit: 1,
+    });
+    request
+        .metadata_mut()
+        .insert("authorization", "Bearer 123".parse().unwrap());
+    let resp = client.search_points(request).await;
+
+    // check if request is successful
+    assert!(resp.is_ok());
+    let result = resp.unwrap().into_inner();
+
+    // 1 vector has to be returned
+    assert_eq!(result.result_point_ids.len(), 1);
+
+    // check if the returned point id matches the inserted point id
+    assert_eq!(result.result_point_ids[0], PointId { id: point.id });
+}
 
 #[tokio::test]
-async fn test_unauthorized_rpc() {}
+async fn test_unauthorized_rpc() {
+    start_test_server().await.unwrap();
+    let mut client = create_test_client().await.unwrap();
+
+    // insert a test vector
+    let test_vec = vec![1.0, 2.0, 3.0];
+    let mut request = tonic::Request::new(InsertVectorRequest {
+        vector: Some(DenseVector {
+            values: test_vec.clone(),
+        }),
+        payload: Some(Struct::default()),
+    });
+    request
+        .metadata_mut()
+        .insert("authorization", "Bearer 43121".parse().unwrap()); // wrong auth token
+    let resp = client.insert_vector(request).await;
+
+    // request must fail
+    assert!(resp.is_err());
+}