refactor: associate Stream::commit_schema

Author: Devdutt Shenoi

src/event/mod.rs

@@ -17,22 +17,17 @@
 *
 */
 
-pub mod format;
+use std::{collections::HashMap, sync::Arc};
 
 use arrow_array::RecordBatch;
-use arrow_schema::{Field, Fields, Schema};
-use itertools::Itertools;
-use std::sync::Arc;
-
-use self::error::EventError;
-use crate::{
-    metadata::update_stats,
-    parseable::{StagingError, PARSEABLE},
-    storage::StreamType,
-    LOCK_EXPECT,
-};
+use arrow_schema::Field;
 use chrono::NaiveDateTime;
-use std::collections::HashMap;
+use error::EventError;
+use itertools::Itertools;
+
+use crate::{metadata::update_stats, parseable::PARSEABLE, storage::StreamType};
+
+pub mod format;
 
 pub const DEFAULT_TIMESTAMP_KEY: &str = "p_timestamp";
 pub const USER_AGENT_KEY: &str = "p_user_agent";
@@ -67,11 +62,12 @@ impl Event {
             }
         }
 
+        let stream = PARSEABLE.get_or_create_stream(&self.stream_name);
         if self.is_first_event {
-            commit_schema(&self.stream_name, self.rb.schema())?;
+            stream.commit_schema(self.rb.schema())?;
         }
 
-        PARSEABLE.get_or_create_stream(&self.stream_name).push(
+        stream.push(
             &key,
             &self.rb,
             self.parsed_timestamp,
@@ -117,32 +113,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("ObjectStorage Error: {0}")]
         ObjectStorage(#[from] ObjectStorageError),
+        #[error("{0}")]
+        NotFound(#[from] StreamNotFound),
     }
 }

src/handlers/http/query.rs

@@ -35,8 +35,6 @@ use tracing::error;
 
 use crate::event::error::EventError;
 use crate::handlers::http::fetch_schema;
-
-use crate::event::commit_schema;
 use crate::metrics::QUERY_EXECUTE_TIME;
 use crate::option::Mode;
 use crate::parseable::{StreamNotFound, PARSEABLE};
@@ -178,7 +176,9 @@ pub async fn update_schema_when_distributed(tables: &Vec<String>) -> Result<(),
                     .commit_schema(table, new_schema.clone())
                     .await?;
 
-                commit_schema(table, Arc::new(new_schema))?;
+                PARSEABLE
+                    .get_stream(table)?
+                    .commit_schema(Arc::new(new_schema))?;
             }
         }
     }

src/parseable/streams.rs

@@ -28,7 +28,7 @@ use std::{
 };
 
 use arrow_array::RecordBatch;
-use arrow_schema::{Field, Fields, Schema};
+use arrow_schema::{Field, Fields, Schema, SchemaRef};
 use chrono::{NaiveDateTime, Timelike, Utc};
 use derive_more::{Deref, DerefMut};
 use itertools::Itertools;
@@ -591,6 +591,18 @@ impl Stream {
         self.metadata.read().expect(LOCK_EXPECT).schema_version
     }
 
+    /// Stores updated schema in-memory
+    pub fn commit_schema(&self, schema: SchemaRef) -> 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.as_ref().clone()])?;
+        metadata.schema.clear();
+        metadata
+            .schema
+            .extend(schema.fields.iter().map(|f| (f.name().clone(), f.clone())));
+        Ok(())
+    }
+
     pub fn get_schema(&self) -> Arc<Schema> {
         let metadata = self.metadata.read().expect(LOCK_EXPECT);