GETKEY implementation

Author: deven96

ahnlich/server/src/engine/store.rs

@@ -172,7 +172,7 @@ impl StoreHandler {
         keys: Vec<StoreKey>,
     ) -> Result<Vec<(StoreKey, StoreValue)>, ServerError> {
         let store = self.get(store_name)?;
-        Ok(store.get_keys(keys))
+        store.get_keys(keys)
     }
 
     /// Matches SET - adds new entries into a particular store
@@ -279,12 +279,27 @@ impl Store {
     }
 
     /// Gets a bunch of store keys from the store
-    fn get_keys(&self, val: Vec<StoreKey>) -> Vec<(StoreKey, StoreValue)> {
+    fn get_keys(&self, val: Vec<StoreKey>) -> Result<Vec<(StoreKey, StoreValue)>, ServerError> {
         if val.is_empty() {
-            return vec![];
+            return Ok(vec![]);
         }
-        let keys = val.iter().map(From::from);
-        self.get(keys)
+        // return error if dimensions do not match
+        let keys = val
+            .into_iter()
+            .map(|key| {
+                let store_dimension = self.dimension.get();
+                let input_dimension = key.dimension();
+                if input_dimension != store_dimension {
+                    return Err(ServerError::StoreDimensionMismatch {
+                        store_dimension,
+                        input_dimension,
+                    });
+                }
+                Ok(key)
+            })
+            .collect::<Result<Vec<StoreKey>, ServerError>>()?;
+        let keys = keys.iter().map(From::from);
+        Ok(self.get(keys))
     }
 
     /// Gets a bunch of store entries that matches a predicate condition

ahnlich/server/src/lib.rs

@@ -208,6 +208,11 @@ impl ServerTask {
                     .set_in_store(&store, inputs)
                     .map(ServerResponse::Set)
                     .map_err(|e| format!("{e}")),
+                Query::GetKey { store, keys } => self
+                    .store_handler
+                    .get_key_in_store(&store, keys)
+                    .map(ServerResponse::GetKey)
+                    .map_err(|e| format!("{e}")),
                 _ => Err("Response not implemented".to_string()),
             })
         }

ahnlich/server/tests/server_test.rs

@@ -202,6 +202,102 @@ async fn test_set_in_store() {
     query_server_assert_result(&mut reader, message, expected).await
 }
 
+#[tokio::test]
+async fn test_get_key() {
+    let server = server::Server::new(&ServerConfig::default())
+        .await
+        .expect("Could not initialize server");
+    let address = server.local_addr().expect("Could not get local addr");
+    let _ = tokio::spawn(async move { server.start().await });
+    // Allow some time for the server to start
+    tokio::time::sleep(Duration::from_millis(100)).await;
+    let message = ServerQuery::from_queries(&[
+        // should error as store does not yet exist
+        Query::GetKey {
+            store: StoreName("Main".to_string()),
+            keys: vec![],
+        },
+        Query::CreateStore {
+            store: StoreName("Main".to_string()),
+            dimension: NonZeroUsize::new(2).unwrap(),
+            create_predicates: HashSet::new(),
+            error_if_exists: true,
+        },
+        Query::Set {
+            store: StoreName("Main".to_string()),
+            inputs: vec![
+                (
+                    StoreKey(array![1.0, 0.2]),
+                    HashMap::from_iter([(
+                        MetadataKey::new("title".into()),
+                        MetadataValue::new("sorcerer".into()),
+                    )]),
+                ),
+                (
+                    StoreKey(array![1.2, 0.3]),
+                    HashMap::from_iter([(
+                        MetadataKey::new("title".into()),
+                        MetadataValue::new("elf".into()),
+                    )]),
+                ),
+            ],
+        },
+        // should error as dimension mismatch
+        Query::GetKey {
+            store: StoreName("Main".to_string()),
+            keys: vec![
+                StoreKey(array![0.2, 0.3, 0.4]),
+                StoreKey(array![0.2, 0.3, 0.4, 0.6]),
+                StoreKey(array![0.4, 0.6]),
+            ],
+        },
+        // should not error however the keys do not exist so should be empty
+        Query::GetKey {
+            store: StoreName("Main".to_string()),
+            keys: vec![StoreKey(array![0.4, 0.6]), StoreKey(array![0.2, 0.5])],
+        },
+        // Gets back the existing key in order
+        Query::GetKey {
+            store: StoreName("Main".to_string()),
+            keys: vec![
+                StoreKey(array![1.2, 0.3]),
+                StoreKey(array![0.4, 0.6]),
+                StoreKey(array![1.0, 0.2]),
+            ],
+        },
+    ]);
+    let mut expected = ServerResult::with_capacity(6);
+    expected.push(Err("Store Main not found".to_string()));
+    expected.push(Ok(ServerResponse::Unit));
+    expected.push(Ok(ServerResponse::Set(StoreUpsert {
+        inserted: 2,
+        updated: 0,
+    })));
+    expected.push(Err(
+        "Store dimension is [2], input dimension of [3] was specified".to_string(),
+    ));
+    expected.push(Ok(ServerResponse::GetKey(vec![])));
+    expected.push(Ok(ServerResponse::GetKey(vec![
+        (
+            StoreKey(array![1.2, 0.3]),
+            HashMap::from_iter([(
+                MetadataKey::new("title".into()),
+                MetadataValue::new("elf".into()),
+            )]),
+        ),
+        (
+            StoreKey(array![1.0, 0.2]),
+            HashMap::from_iter([(
+                MetadataKey::new("title".into()),
+                MetadataValue::new("sorcerer".into()),
+            )]),
+        ),
+    ])));
+    let stream = TcpStream::connect(address).await.unwrap();
+    let mut reader = BufReader::new(stream);
+    query_server_assert_result(&mut reader, message, expected).await
+}
+
 #[tokio::test]
 async fn test_drop_stores() {
     let server = server::Server::new(&ServerConfig::default())

ahnlich/types/src/keyval.rs

@@ -25,6 +25,12 @@ pub type StoreValue = StdHashMap<MetadataKey, MetadataValue>;
 #[serde(transparent)]
 pub struct StoreKey(pub Array1<f64>);
 
+impl StoreKey {
+    pub fn dimension(&self) -> usize {
+        self.0.len()
+    }
+}
+
 impl Eq for StoreKey {}
 
 impl PartialEq for StoreKey {

ahnlich/types/src/query.rs

@@ -25,6 +25,7 @@ pub enum Query {
         error_if_exists: bool,
     },
     GetKey {
+        store: StoreName,
         keys: Vec<StoreKey>,
     },
     GetPred {

ahnlich/types/src/server.rs

@@ -1,5 +1,7 @@
 use crate::bincode::BinCodeSerAndDeser;
+use crate::keyval::StoreKey;
 use crate::keyval::StoreName;
+use crate::keyval::StoreValue;
 use serde::Deserialize;
 use serde::Serialize;
 use std::collections::HashSet;
@@ -17,6 +19,7 @@ pub enum ServerResponse {
     StoreList(HashSet<StoreInfo>),
     InfoServer(ServerInfo),
     Set(StoreUpsert),
+    GetKey(Vec<(StoreKey, StoreValue)>),
     // TODO: Define return types for queries, e.t.c
 }