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

@@ -182,14 +182,14 @@ impl AlertTrait for ThresholdAlert {
         &mut self,
         new_notification_state: NotificationState,
     ) -> Result<(), AlertError> {
+        // update state in memory
+        self.notification_state = new_notification_state;
+
         // update on disk
         PARSEABLE
             .metastore
             .put_alert(&self.to_alert_config())
             .await?;
-        // update state in memory
-        self.notification_state = new_notification_state;
-
         Ok(())
     }
 

src/alerts/alert_types.rs

@@ -992,11 +992,11 @@ impl actix_web::ResponseError for AlertError {
 impl AlertManagerTrait for Alerts {
     /// Loads alerts from disk, blocks
     async fn load(&self) -> anyhow::Result<()> {
-        let mut map = self.alerts.write().await;
-
         // Get alerts path and read raw bytes for migration handling
         let raw_objects = PARSEABLE.metastore.get_alerts().await?;
 
+        let mut map = self.alerts.write().await;
+
         for raw_bytes in raw_objects {
             // First, try to parse as JSON Value to check version
             let json_value: JsonValue = match serde_json::from_slice(&raw_bytes) {

src/alerts/mod.rs

@@ -90,8 +90,7 @@ pub async fn list(req: HttpRequest) -> Result<impl Responder, StreamError> {
     let res = PARSEABLE
         .metastore
         .list_streams()
-        .await
-        .unwrap()
+        .await?
         .into_iter()
         .filter(|logstream| {
             Users.authorize(key.clone(), Action::ListStream, Some(logstream), None)

src/handlers/http/logstream.rs

@@ -163,7 +163,6 @@ pub async fn get_stats(
         })?;
 
         if !date_value.is_empty() {
-            // this function requires all the ingestor stream jsons
             let obs = PARSEABLE
                 .metastore
                 .get_all_stream_jsons(&stream_name, None)

src/handlers/http/modal/query/querier_logstream.rs

@@ -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> {