Schema Support - Logical Separation for Stores

ahnlich/ai/src/engine/mod.rs

@@ -1,3 +1,4 @@
 pub mod ai;
 pub mod operations;
 pub mod store;
+pub mod versioned;

ahnlich/ai/src/engine/operations.rs

@@ -1,6 +1,7 @@
 use std::{collections::HashMap, sync::Arc};
 
 use ahnlich_client_rs::db::DbClient;
+use ahnlich_types::ai::query::DropSchema as AiDropSchema;
 use ahnlich_types::{
     ai::{
         models::AiModel,
@@ -28,6 +29,7 @@ use ahnlich_types::{
         In, Predicate, PredicateCondition, predicate::Kind as PredicateKind,
         predicate_condition::Kind,
     },
+    schema::Schema,
 };
 use itertools::Itertools;
 use rayon::prelude::*;
@@ -90,6 +92,13 @@ pub async fn create_store(
         predicates.push(AHNLICH_AI_ONE_TO_MANY_INDEX_META_KEY.to_string());
     }
 
+    let schema = match &params.schema {
+        Some(s) => {
+            Schema::try_new(s.clone()).map_err(|e| AIProxyError::InvalidArgument(e.to_owned()))?
+        }
+        None => Schema::default(),
+    };
+
     require_db_client(db_client)?
         .create_store(
             DbCreateStore {
@@ -98,6 +107,7 @@ pub async fn create_store(
                 create_predicates: predicates,
                 non_linear_indices: params.non_linear_indices,
                 error_if_exists: params.error_if_exists,
+                schema: Some(schema.to_string()),
             },
             parent_id,
         )
@@ -107,6 +117,7 @@ pub async fn create_store(
         StoreName {
             value: params.store,
         },
+        &schema,
         query_model,
         index_model,
         params.error_if_exists,
@@ -353,6 +364,12 @@ pub async fn drop_store(
     let store_name = StoreName {
         value: params.store,
     };
+    let schema = match &params.schema {
+        Some(s) => {
+            Schema::try_new(s.clone()).map_err(|e| AIProxyError::InvalidArgument(e.to_owned()))?
+        }
+        None => Schema::default(),
+    };
 
     if store_handler.get(&store_name).is_ok()
         && let Some(db_client) = db_client
@@ -362,13 +379,14 @@ pub async fn drop_store(
                 DbDropStore {
                     store: store_name.value.clone(),
                     error_if_not_exists: params.error_if_not_exists,
+                    schema: Some(schema.to_string()),
                 },
                 parent_id,
             )
             .await?;
     }
 
-    let dropped = store_handler.drop_store(store_name, params.error_if_not_exists)?;
+    let dropped = store_handler.drop_store(store_name, &schema, params.error_if_not_exists)?;
     Ok(Del {
         deleted_count: dropped as u64,
     })
@@ -380,19 +398,14 @@ pub async fn purge_stores(
     _params: PurgeStores,
     parent_id: Option<String>,
 ) -> Result<Del, Status> {
-    let store_names: Vec<StoreName> = store_handler
-        .list_stores()
-        .into_iter()
-        .map(|store| StoreName { value: store.name })
-        .collect();
-
     if let Some(db_client) = db_client {
-        for store_name in &store_names {
+        for (schema, store_name) in store_handler.all_store_names_by_schema() {
             db_client
                 .drop_store(
                     DbDropStore {
-                        store: store_name.value.clone(),
+                        store: store_name.value,
                         error_if_not_exists: false,
+                        schema: Some(schema),
                     },
                     parent_id.clone(),
                 )
@@ -408,10 +421,19 @@ pub async fn purge_stores(
 pub async fn list_stores(
     store_handler: &AIStoreHandler,
     db_client: Option<Arc<DbClient>>,
-    _params: ListStores,
+    params: ListStores,
     parent_id: Option<String>,
 ) -> Result<StoreList, Status> {
-    let mut stores: Vec<AiStoreInfo> = store_handler.list_stores().into_iter().sorted().collect();
+    let filter_schema = params
+        .schema
+        .as_ref()
+        .filter(|s| !s.is_empty())
+        .map(|s| Schema::new(s.clone()));
+    let mut stores: Vec<AiStoreInfo> = store_handler
+        .list_stores(filter_schema.as_ref())
+        .into_iter()
+        .sorted()
+        .collect();
 
     if let Some(db_client) = db_client
         && let Ok(db_store_list) = db_client.list_stores(parent_id).await
@@ -439,3 +461,18 @@ pub async fn list_stores(
 
     Ok(StoreList { stores })
 }
+
+pub async fn drop_schema(
+    store_handler: &AIStoreHandler,
+    _db_client: Option<Arc<DbClient>>,
+    params: AiDropSchema,
+    _parent_id: Option<String>,
+) -> Result<Del, Status> {
+    let schema =
+        Schema::try_new(params.schema).map_err(|e| AIProxyError::InvalidArgument(e.to_owned()))?;
+
+    let dropped = store_handler.drop_schema(&schema)?;
+    Ok(Del {
+        deleted_count: dropped as u64,
+    })
+}