Lookup store refactoring

5 files changed, 95 insertions, 98 deletions

diff --git a/crates/store/src/backend/memory/lookup.rs b/crates/store/src/backend/memory/lookup.rs
index 41a4dc08..4538fd21 100644
--- a/crates/store/src/backend/memory/lookup.rs
+++ b/crates/store/src/backend/memory/lookup.rs
@@ -21,57 +21,9 @@
  * for more details.
 */
 
-use crate::{IntoRows, QueryResult, QueryType, Row, Value};
+use crate::{IntoRows, Row};
 
-use super::{LookupList, MatchType, MemoryStore};
-
-impl MemoryStore {
-    pub(crate) fn query<T: QueryResult>(
-        &self,
-        _: &str,
-        params: Vec<Value<'_>>,
-    ) -> crate::Result<T> {
-        let exists = match T::query_type() {
-            QueryType::Exists => true,
-            QueryType::QueryOne => false,
-            QueryType::QueryAll | QueryType::Execute => {
-                return Err(crate::Error::InternalError(
-                    "Unsupported query type".to_string(),
-                ))
-            }
-        };
-
-        let needle = params.first().map(|v| v.to_str()).unwrap_or_default();
-
-        match self {
-            MemoryStore::List(list) => {
-                let found = list.contains(needle.as_ref());
-                if exists {
-                    Ok(T::from_exists(found))
-                } else {
-                    Ok(T::from_query_one(Some(Row {
-                        values: vec![Value::Bool(found)],
-                    })))
-                }
-            }
-            MemoryStore::Map(map) => {
-                if let Some(value) = map.get(needle.as_ref()) {
-                    if exists {
-                        Ok(T::from_exists(true))
-                    } else {
-                        Ok(T::from_query_one(Some(Row {
-                            values: vec![value.clone()],
-                        })))
-                    }
-                } else if exists {
-                    Ok(T::from_exists(false))
-                } else {
-                    Ok(T::from_query_one(None::<Row>))
-                }
-            }
-        }
-    }
-}
+use super::{LookupList, MatchType};
 
 impl IntoRows for Option<Row> {
     fn into_row(self) -> Option<Row> {
diff --git a/crates/store/src/backend/memory/main.rs b/crates/store/src/backend/memory/main.rs
index 9274331f..ec56b2d3 100644
--- a/crates/store/src/backend/memory/main.rs
+++ b/crates/store/src/backend/memory/main.rs
@@ -54,7 +54,7 @@ impl MemoryStore {
     pub async fn open(config: &Config, prefix: impl AsKey) -> crate::Result<Self> {
         let prefix = prefix.as_key();
 
-        let lookup_type = config.property_require::<LookupType>((&prefix, "type"))?;
+        let lookup_type = config.property_require::<LookupType>((&prefix, "format"))?;
         let format = LookupFormat {
             lookup_type,
             comment: config.value((&prefix, "comment")).map(|s| s.to_string()),
diff --git a/crates/store/src/config.rs b/crates/store/src/config.rs
index 355e4106..7d66a3e4 100644
--- a/crates/store/src/config.rs
+++ b/crates/store/src/config.rs
@@ -24,12 +24,12 @@
 use std::sync::Arc;
 
 use async_trait::async_trait;
-use utils::config::{cron::SimpleCron, utils::AsKey, Config};
+use utils::config::{cron::SimpleCron, Config};
 
 use crate::{
     backend::{fs::FsStore, memory::MemoryStore},
     write::purge::{PurgeSchedule, PurgeStore},
-    Lookup, LookupStore, Store, Stores,
+    LookupStore, QueryStore, Store, Stores,
 };
 
 #[cfg(feature = "s3")]
@@ -178,21 +178,9 @@ impl ConfigStore for Config {
                     continue;
                 }
                 "memory" => {
-                    let prefix = prefix.as_key();
-                    for lookup_id in self.sub_keys((&prefix, "lookup")) {
-                        config.lookups.insert(
-                            format!("{store_id}/{lookup_id}"),
-                            Arc::new(Lookup {
-                                store: MemoryStore::open(
-                                    self,
-                                    (prefix.as_str(), "lookup", lookup_id),
-                                )
-                                .await?
-                                .into(),
-                                query: String::new(),
-                            }),
-                        );
-                    }
+                    config
+                        .lookup_stores
+                        .insert(store_id, MemoryStore::open(self, prefix).await?.into());
                     continue;
                 }
 
@@ -202,15 +190,15 @@ impl ConfigStore for Config {
                 }
             };
 
-            // Add queries
+            // Add queries as lookup stores
             let lookup_store: LookupStore = lookup_store.into();
             for lookup_id in self.sub_keys(("store", id, "query")) {
-                config.lookups.insert(
+                config.lookup_stores.insert(
                     format!("{store_id}/{lookup_id}"),
-                    Arc::new(Lookup {
+                    LookupStore::Query(Arc::new(QueryStore {
                         store: lookup_store.clone(),
                         query: self.property_require(("store", id, "query", lookup_id))?,
-                    }),
+                    })),
                 );
             }
             config.lookup_stores.insert(store_id, lookup_store.clone());
diff --git a/crates/store/src/dispatch/lookup.rs b/crates/store/src/dispatch/lookup.rs
index 580f9f43..a1893778 100644
--- a/crates/store/src/dispatch/lookup.rs
+++ b/crates/store/src/dispatch/lookup.rs
@@ -21,6 +21,7 @@
  * for more details.
 */
 
+use crate::{backend::memory::MemoryStore, Row};
 #[allow(unused_imports)]
 use crate::{
     write::{
@@ -39,21 +40,14 @@ impl LookupStore {
         params: Vec<Value<'_>>,
     ) -> crate::Result<T> {
         let result = match self {
-            LookupStore::Store(store) => match store {
-                #[cfg(feature = "sqlite")]
-                Store::SQLite(store) => store.query(query, params).await,
-                #[cfg(feature = "postgres")]
-                Store::PostgreSQL(store) => store.query(query, params).await,
-                #[cfg(feature = "mysql")]
-                Store::MySQL(store) => store.query(query, params).await,
-                _ => Err(crate::Error::InternalError(
-                    "Store does not support queries".into(),
-                )),
-            },
-            LookupStore::Memory(store) => store.query(query, params),
-            #[cfg(feature = "redis")]
-            LookupStore::Redis(_) => Err(crate::Error::InternalError(
-                "Redis does not support queries".into(),
+            #[cfg(feature = "sqlite")]
+            LookupStore::Store(Store::SQLite(store)) => store.query(query, params).await,
+            #[cfg(feature = "postgres")]
+            LookupStore::Store(Store::PostgreSQL(store)) => store.query(query, params).await,
+            #[cfg(feature = "mysql")]
+            LookupStore::Store(Store::MySQL(store)) => store.query(query, params).await,
+            _ => Err(crate::Error::InternalError(
+                "Store does not support queries".into(),
             )),
         };
 
@@ -89,11 +83,21 @@ impl LookupStore {
             }
             #[cfg(feature = "redis")]
             LookupStore::Redis(store) => store.key_set(key, value).await,
-            LookupStore::Memory(_) => unimplemented!(),
+            LookupStore::Query(lookup) => lookup
+                .store
+                .query::<usize>(
+                    &lookup.query,
+                    vec![String::from_utf8(key).unwrap_or_default().into()],
+                )
+                .await
+                .map(|_| ()),
+            LookupStore::Memory(_) => Err(crate::Error::InternalError(
+                "This store does not support key_set".into(),
+            )),
         }
     }
 
-    pub async fn key_get<T: Deserialize + std::fmt::Debug + 'static>(
+    pub async fn key_get<T: Deserialize + From<Value<'static>> + std::fmt::Debug + 'static>(
         &self,
         key: LookupKey,
     ) -> crate::Result<LookupValue<T>> {
@@ -120,7 +124,38 @@ impl LookupStore {
             },
             #[cfg(feature = "redis")]
             LookupStore::Redis(store) => store.key_get(key).await,
-            LookupStore::Memory(_) => unimplemented!(),
+            LookupStore::Memory(store) => {
+                let key = String::from(key);
+                match store.as_ref() {
+                    MemoryStore::List(list) => Ok(if list.contains(&key) {
+                        LookupValue::Value {
+                            value: T::from(Value::Bool(true)),
+                            expires: 0,
+                        }
+                    } else {
+                        LookupValue::None
+                    }),
+                    MemoryStore::Map(map) => Ok(map
+                        .get(&key)
+                        .map(|value| LookupValue::Value {
+                            value: T::from(value.to_owned()),
+                            expires: 0,
+                        })
+                        .unwrap_or(LookupValue::None)),
+                }
+            }
+            LookupStore::Query(lookup) => lookup
+                .store
+                .query::<Option<Row>>(&lookup.query, vec![String::from(key).into()])
+                .await
+                .map(|row| {
+                    row.and_then(|row| row.values.into_iter().next())
+                        .map(|value| LookupValue::Value {
+                            value: T::from(value),
+                            expires: 0,
+                        })
+                        .unwrap_or(LookupValue::None)
+                }),
         }
     }
 
@@ -169,7 +204,7 @@ impl LookupStore {
             }
             #[cfg(feature = "redis")]
             LookupStore::Redis(_) => {}
-            LookupStore::Memory(_) => {}
+            LookupStore::Memory(_) | LookupStore::Query(_) => {}
         }
 
         Ok(())
@@ -190,3 +225,16 @@ impl<T: Deserialize> Deserialize for LookupValue<T> {
         })
     }
 }
+
+impl From<Value<'static>> for String {
+    fn from(value: Value<'static>) -> Self {
+        match value {
+            Value::Text(string) => string.into_owned(),
+            Value::Blob(bytes) => String::from_utf8_lossy(bytes.as_ref()).into_owned(),
+            Value::Bool(boolean) => boolean.to_string(),
+            Value::Null => String::new(),
+            Value::Integer(num) => num.to_string(),
+            Value::Float(num) => num.to_string(),
+        }
+    }
+}
diff --git a/crates/store/src/lib.rs b/crates/store/src/lib.rs
index f780b6fb..06f2aca0 100644
--- a/crates/store/src/lib.rs
+++ b/crates/store/src/lib.rs
@@ -194,13 +194,6 @@ pub struct Stores {
     pub blob_stores: AHashMap<String, BlobStore>,
     pub fts_stores: AHashMap<String, FtsStore>,
     pub lookup_stores: AHashMap<String, LookupStore>,
-    pub lookups: AHashMap<String, Arc<Lookup>>,
-}
-
-#[derive(Clone)]
-pub struct Lookup {
-    pub store: LookupStore,
-    pub query: String,
 }
 
 #[derive(Clone)]
@@ -235,11 +228,17 @@ pub enum FtsStore {
 #[derive(Clone)]
 pub enum LookupStore {
     Store(Store),
+    Query(Arc<QueryStore>),
     Memory(Arc<MemoryStore>),
     #[cfg(feature = "redis")]
     Redis(Arc<RedisStore>),
 }
 
+pub struct QueryStore {
+    pub store: LookupStore,
+    pub query: String,
+}
+
 #[cfg(feature = "sqlite")]
 impl From<SqliteStore> for Store {
     fn from(store: SqliteStore) -> Self {
@@ -364,6 +363,16 @@ impl<'x> Value<'x> {
     }
 }
 
+impl From<LookupKey> for String {
+    fn from(value: LookupKey) -> Self {
+        let key = match value {
+            LookupKey::Key(key) | LookupKey::Counter(key) => key,
+        };
+        String::from_utf8(key)
+            .unwrap_or_else(|err| String::from_utf8_lossy(&err.into_bytes()).into_owned())
+    }
+}
+
 #[derive(Clone, Debug)]
 pub struct Row {
     pub values: Vec<Value<'static>>,