refactor: Parseable::commit_schema

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, Stream, PARSEABLE},
-    storage::StreamType,
-    LOCK_EXPECT,
-};
+use crate::{metadata::update_stats, parseable::Stream, storage::StreamType};
 use chrono::NaiveDateTime;
 use std::collections::HashMap;
 
@@ -71,7 +66,8 @@ impl Event {
             }
 
             if self.is_first_event {
-                commit_schema(&stream.stream_name, partition.rb.schema())?;
+                let schema = partition.rb.schema().as_ref().clone();
+                stream.commit_schema(schema)?;
             }
 
             stream.push(
@@ -122,23 +118,6 @@ 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};

src/handlers/http/query.rs

@@ -29,14 +29,12 @@ use serde::{Deserialize, Serialize};
 use serde_json::{json, Value};
 use std::collections::HashMap;
 use std::pin::Pin;
-use std::sync::Arc;
 use std::time::Instant;
 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::PARSEABLE;
@@ -174,7 +172,9 @@ pub async fn update_schema_when_distributed(tables: &Vec<String>) -> Result<(),
                 // commit schema merges the schema internally and updates the schema in storage.
                 commit_schema_to_storage(table, new_schema.clone()).await?;
 
-                commit_schema(table, Arc::new(new_schema))?;
+                PARSEABLE
+                    .get_or_create_stream(table)
+                    .commit_schema(new_schema)?;
             }
         }
     }

src/parseable/streams.rs

@@ -624,6 +624,19 @@ impl Stream {
         Arc::new(Schema::new(fields))
     }
 
+    pub fn commit_schema(&self, schema: Schema) -> Result<(), StagingError> {
+        let current_schema = self.get_schema().as_ref().clone();
+        let updated_schema = Schema::try_merge([current_schema, schema])?
+            .fields
+            .into_iter()
+            .map(|field| (field.name().to_owned(), field.clone()))
+            .collect();
+
+        self.metadata.write().expect(LOCK_EXPECT).schema = updated_schema;
+
+        Ok(())
+    }
+
     pub fn get_schema_raw(&self) -> HashMap<String, Arc<Field>> {
         self.metadata.read().expect(LOCK_EXPECT).schema.clone()
     }