Add userchats to metastore

Author: parmesant

Cargo.lock

@@ -142,6 +142,7 @@ futures-core = "0.3.31"
 tempfile = "3.20.0"
 lazy_static = "1.4.0"
 prost = "0.13.1"
+dashmap = "6.1.0"
 
 [build-dependencies]
 cargo_toml = "0.21"

Cargo.toml

@@ -21,6 +21,7 @@ use std::collections::{BTreeMap, HashSet};
 use arrow_schema::Schema;
 use bytes::Bytes;
 use chrono::{DateTime, Utc};
+use dashmap::DashMap;
 use erased_serde::Serialize as ErasedSerialize;
 use tonic::async_trait;
 
@@ -58,6 +59,11 @@ pub trait Metastore: std::fmt::Debug + Send + Sync {
     async fn put_dashboard(&self, obj: &dyn MetastoreObject) -> Result<(), MetastoreError>;
     async fn delete_dashboard(&self, obj: &dyn MetastoreObject) -> Result<(), MetastoreError>;
 
+    /// chats
+    async fn get_chats(&self) -> Result<DashMap<String, Vec<Bytes>>, MetastoreError>;
+    async fn put_chat(&self, obj: &dyn MetastoreObject) -> Result<(), MetastoreError>;
+    async fn delete_chat(&self, obj: &dyn MetastoreObject) -> Result<(), MetastoreError>;
+
     /// filters
     async fn get_filters(&self) -> Result<Vec<Filter>, MetastoreError>;
     async fn put_filter(&self, obj: &dyn MetastoreObject) -> Result<(), MetastoreError>;

src/metastore/metastore_traits.rs

@@ -24,6 +24,7 @@ use std::{
 use arrow_schema::Schema;
 use bytes::Bytes;
 use chrono::{DateTime, Utc};
+use dashmap::DashMap;
 use http::StatusCode;
 use relative_path::RelativePathBuf;
 use tonic::async_trait;
@@ -188,6 +189,54 @@ impl Metastore for ObjectStoreMetastore {
             .await?)
     }
 
+    /// Fetch all chats
+    async fn get_chats(&self) -> Result<DashMap<String, Vec<Bytes>>, MetastoreError> {
+        let all_user_chats = DashMap::new();
+
+        let users_dir = RelativePathBuf::from(USERS_ROOT_DIR);
+        for user in self.storage.list_dirs_relative(&users_dir).await? {
+            if user.starts_with(".") {
+                continue;
+            }
+            let mut chats = Vec::new();
+            let chats_path = users_dir.join(&user).join("chats");
+            let user_chats = self
+                .storage
+                .get_objects(
+                    Some(&chats_path),
+                    Box::new(|file_name| file_name.ends_with(".json")),
+                )
+                .await?;
+            for chat in user_chats {
+                chats.push(chat);
+            }
+
+            all_user_chats.insert(user, chats);
+        }
+
+        Ok(all_user_chats)
+    }
+
+    /// Save a chat
+    async fn put_chat(&self, obj: &dyn MetastoreObject) -> Result<(), MetastoreError> {
+        // we need the path to store in obj store
+        let path = obj.get_object_path();
+
+        Ok(self
+            .storage
+            .put_object(&RelativePathBuf::from(path), to_bytes(obj))
+            .await?)
+    }
+
+    /// Delete a chat
+    async fn delete_chat(&self, obj: &dyn MetastoreObject) -> Result<(), MetastoreError> {
+        let path = obj.get_object_path();
+        Ok(self
+            .storage
+            .delete_object(&RelativePathBuf::from(path))
+            .await?)
+    }
+
     // for get filters, take care of migration and removal of incorrect/old filters
     // return deserialized filter
     async fn get_filters(&self) -> Result<Vec<Filter>, MetastoreError> {