refactor: associate ObjectStorage::fetch_schemas

Author: Devdutt Shenoi

src/event/mod.rs

@@ -20,17 +20,12 @@
 pub mod format;
 
 use arrow_array::RecordBatch;
-use arrow_schema::{Field, Fields, Schema};
+use arrow_schema::Field;
 use itertools::Itertools;
 use std::sync::Arc;
 
 use self::error::EventError;
-use crate::{
-    metadata::update_stats,
-    parseable::{StagingError, PARSEABLE},
-    storage::StreamType,
-    LOCK_EXPECT,
-};
+use crate::{metadata::update_stats, parseable::PARSEABLE, storage::StreamType};
 use chrono::NaiveDateTime;
 use std::collections::HashMap;
 
@@ -64,11 +59,13 @@ impl Event {
             }
         }
 
+        let stream = PARSEABLE.get_stream(&self.stream_name)?;
         if self.is_first_event {
-            commit_schema(&self.stream_name, self.rb.schema())?;
+            let schema = self.rb.schema().as_ref().clone();
+            stream.commit_schema(schema)?;
         }
 
-        PARSEABLE.get_or_create_stream(&self.stream_name).push(
+        stream.push(
             &key,
             &self.rb,
             self.parsed_timestamp,
@@ -114,31 +111,20 @@ pub fn get_schema_key(fields: &[Arc<Field>]) -> String {
     format!("{hash:x}")
 }
 
-pub fn commit_schema(stream_name: &str, schema: Arc<Schema>) -> Result<(), StagingError> {
-    let mut stream_metadata = PARSEABLE.streams.write().expect("lock poisoned");
-
-    let map = &mut stream_metadata
-        .get_mut(stream_name)
-        .expect("map has entry for this stream name")
-        .metadata
-        .write()
-        .expect(LOCK_EXPECT)
-        .schema;
-    let current_schema = Schema::new(map.values().cloned().collect::<Fields>());
-    let schema = Schema::try_merge(vec![current_schema, schema.as_ref().clone()])?;
-    map.clear();
-    map.extend(schema.fields.iter().map(|f| (f.name().clone(), f.clone())));
-    Ok(())
-}
-
 pub mod error {
-
-    use crate::{parseable::StagingError, storage::ObjectStorageError};
+    use crate::{
+        parseable::{StagingError, StreamNotFound},
+        storage::ObjectStorageError,
+    };
 
     #[derive(Debug, thiserror::Error)]
     pub enum EventError {
         #[error("Staging Failed: {0}")]
         Staging(#[from] StagingError),
+        #[error("{0}")]
+        Stream(#[from] StreamNotFound),
+        #[error("Arrow Error: {0}")]
+        Arrow(#[from] arrow_schema::ArrowError),
         #[error("ObjectStorage Error: {0}")]
         ObjectStorage(#[from] ObjectStorageError),
     }

src/parseable/mod.rs

@@ -25,9 +25,7 @@ use bytes::Bytes;
 use chrono::Utc;
 use clap::{error::ErrorKind, Parser};
 use http::{header::CONTENT_TYPE, HeaderName, HeaderValue, StatusCode};
-use itertools::Itertools;
 use once_cell::sync::Lazy;
-use relative_path::RelativePathBuf;
 pub use staging::StagingError;
 use streams::StreamRef;
 pub use streams::{StreamNotFound, Streams};
@@ -37,7 +35,7 @@ use tracing::error;
 use crate::connectors::kafka::config::KafkaConfig;
 use crate::{
     cli::{Cli, Options, StorageOptions},
-    event::{commit_schema, error::EventError, format::LogSource},
+    event::{error::EventError, format::LogSource},
     handlers::{
         http::{
             cluster::{sync_streams_with_ingestors, INTERNAL_STREAM_NAME},
@@ -52,7 +50,7 @@ use crate::{
     static_schema::{convert_static_schema_to_arrow_schema, StaticSchema},
     storage::{
         object_storage::parseable_json_path, ObjectStorageError, ObjectStorageProvider,
-        ObjectStoreFormat, Owner, Permisssion, StreamType, STREAM_ROOT_DIRECTORY,
+        ObjectStoreFormat, Owner, Permisssion, StreamType,
     },
     validator,
 };
@@ -768,47 +766,21 @@ impl Parseable {
         Some(first_event_at.to_string())
     }
 
-    /// Fetches the schema for the specified stream.
-    ///
-    /// # Arguments
-    ///
-    /// * `stream_name` - The name of the stream to fetch the schema for.
-    ///
-    /// # Returns
-    ///
-    /// An `anyhow::Result` containing the `arrow_schema::Schema` for the specified stream.
-    pub async fn fetch_schema(&self, stream_name: &str) -> anyhow::Result<Schema> {
-        let path_prefix = RelativePathBuf::from_iter([stream_name, STREAM_ROOT_DIRECTORY]);
-        let store = self.storage.get_object_store();
-        let res: Vec<Schema> = store
-            .get_objects(
-                Some(&path_prefix),
-                Box::new(|file_name: String| file_name.contains(".schema")),
-            )
-            .await?
-            .iter()
-            // we should be able to unwrap as we know the data is valid schema
-            .map(|byte_obj| serde_json::from_slice(byte_obj).expect("data is valid json"))
-            .collect_vec();
-
-        let new_schema = Schema::try_merge(res)?;
-        Ok(new_schema)
-    }
-
     pub async fn update_schema_when_distributed(
         &self,
         tables: &Vec<String>,
     ) -> Result<(), EventError> {
         if self.options.mode == Mode::Query {
             for table in tables {
-                if let Ok(new_schema) = self.fetch_schema(table).await {
+                if let Ok(schemas) = self.storage.get_object_store().fetch_schemas(table).await {
+                    let new_schema = Schema::try_merge(schemas)?;
                     // commit schema merges the schema internally and updates the schema in storage.
                     self.storage
                         .get_object_store()
                         .commit_schema(table, new_schema.clone())
                         .await?;
 
-                    commit_schema(table, Arc::new(new_schema))?;
+                    self.get_stream(table)?.commit_schema(new_schema)?;
                 }
             }
         }

src/parseable/streams.rs

@@ -547,6 +547,17 @@ impl Stream {
         Schema::try_merge(vec![schema, current_schema]).unwrap()
     }
 
+    pub fn commit_schema(&self, schema: Schema) -> Result<(), StagingError> {
+        let mut metadata = self.metadata.write().expect(LOCK_EXPECT);
+        let current_schema = Schema::new(metadata.schema.values().cloned().collect::<Fields>());
+        let schema = Schema::try_merge(vec![current_schema, schema])?;
+        metadata.schema.clear();
+        metadata
+            .schema
+            .extend(schema.fields.iter().map(|f| (f.name().clone(), f.clone())));
+        Ok(())
+    }
+
     /// Stores the provided stream metadata in memory mapping
     pub async fn set_metadata(&self, updated_metadata: LogStreamMetadata) {
         *self.metadata.write().expect(LOCK_EXPECT) = updated_metadata;

src/storage/object_storage.rs

@@ -799,6 +799,32 @@ pub trait ObjectStorage: Debug + Send + Sync + 'static {
         let new_schema = Schema::try_merge(vec![schema, stream_schema]).unwrap();
         self.put_schema(stream_name, &new_schema).await
     }
+
+    /// Fetches all schemas associated with a specified stream.
+    ///
+    /// # Arguments
+    ///
+    /// * `stream_name` - The name of the stream to fetch the schema for.
+    ///
+    /// # Returns
+    ///
+    /// An array of `arrow_schema::Schema` for the specified stream.
+    async fn fetch_schemas(&self, stream_name: &str) -> Result<Vec<Schema>, ObjectStorageError> {
+        let path_prefix = RelativePathBuf::from_iter([stream_name, STREAM_ROOT_DIRECTORY]);
+        let mut schemas: Vec<Schema> = vec![];
+        for byte_obj in self
+            .get_objects(
+                Some(&path_prefix),
+                Box::new(|file_name: String| file_name.contains(".schema")),
+            )
+            .await?
+        {
+            let schema = serde_json::from_slice(&byte_obj)?;
+            schemas.push(schema);
+        }
+
+        Ok(schemas)
+    }
 }
 
 #[inline(always)]