refactor: Parseable::update_schema_when_distributed

Author: Devdutt Shenoi

src/handlers/airplane.rs

@@ -34,7 +34,7 @@ use tonic::transport::{Identity, Server, ServerTlsConfig};
 use tonic_web::GrpcWebLayer;
 
 use crate::handlers::http::cluster::get_ingestor_info;
-use crate::handlers::http::query::{into_query, update_schema_when_distributed};
+use crate::handlers::http::query::into_query;
 use crate::handlers::livetail::cross_origin_config;
 use crate::metrics::QUERY_EXECUTE_TIME;
 use crate::parseable::PARSEABLE;
@@ -156,7 +156,8 @@ impl FlightService for AirServiceImpl {
             .ok_or_else(|| Status::aborted("Malformed SQL Provided, Table Name Not Found"))?
             .to_owned();
 
-        update_schema_when_distributed(&streams)
+        PARSEABLE
+            .update_schema_when_distributed(&streams)
             .await
             .map_err(|err| Status::internal(err.to_string()))?;
 

src/handlers/http/logstream.rs

@@ -18,7 +18,6 @@
 
 use self::error::StreamError;
 use super::cluster::utils::{IngestionStats, QueriedStats, StorageStats};
-use super::query::update_schema_when_distributed;
 use crate::event::format::override_data_type;
 use crate::hottier::{HotTierManager, StreamHotTier, CURRENT_HOT_TIER_VERSION};
 use crate::metadata::SchemaVersion;
@@ -129,7 +128,10 @@ pub async fn get_schema(stream_name: Path<String>) -> Result<impl Responder, Str
     }
 
     let stream = PARSEABLE.get_stream(&stream_name)?;
-    match update_schema_when_distributed(&vec![stream_name.clone()]).await {
+    match PARSEABLE
+        .update_schema_when_distributed(&vec![stream_name.clone()])
+        .await
+    {
         Ok(_) => {
             let schema = stream.get_schema();
             Ok((web::Json(schema), StatusCode::OK))

src/handlers/http/query.rs

@@ -29,14 +29,11 @@ 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::commit_schema;
 use crate::event::error::EventError;
 use crate::metrics::QUERY_EXECUTE_TIME;
-use crate::option::Mode;
 use crate::parseable::PARSEABLE;
 use crate::query::error::ExecuteError;
 use crate::query::{CountsRequest, CountsResponse, Query as LogicalQuery};
@@ -86,7 +83,7 @@ pub async fn query(req: HttpRequest, query_request: Query) -> Result<HttpRespons
     let _ = raw_logical_plan.visit(&mut visitor);
 
     let tables = visitor.into_inner();
-    update_schema_when_distributed(&tables).await?;
+    PARSEABLE.update_schema_when_distributed(&tables).await?;
     let query: LogicalQuery = into_query(&query_request, &session_state, time_range).await?;
 
     let creds = extract_session_key_from_req(&req)?;
@@ -164,25 +161,6 @@ pub async fn get_counts(
     }))
 }
 
-pub async fn update_schema_when_distributed(tables: &Vec<String>) -> Result<(), EventError> {
-    if PARSEABLE.options.mode == Mode::Query {
-        for table in tables {
-            if let Ok(new_schema) = PARSEABLE.fetch_schema(table).await {
-                // commit schema merges the schema internally and updates the schema in storage.
-                PARSEABLE
-                    .storage
-                    .get_object_store()
-                    .commit_schema(table, new_schema.clone())
-                    .await?;
-
-                commit_schema(table, Arc::new(new_schema))?;
-            }
-        }
-    }
-
-    Ok(())
-}
-
 /// Create streams for querier if they do not exist
 /// get list of streams from memory and storage
 /// create streams for memory from storage if they do not exist

src/parseable/mod.rs

@@ -37,7 +37,7 @@ use tracing::error;
 use crate::connectors::kafka::config::KafkaConfig;
 use crate::{
     cli::{Cli, Options, StorageOptions},
-    event::format::LogSource,
+    event::{commit_schema, error::EventError, format::LogSource},
     handlers::{
         http::{
             cluster::{sync_streams_with_ingestors, INTERNAL_STREAM_NAME},
@@ -794,6 +794,27 @@ impl Parseable {
         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 {
+                    // 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))?;
+                }
+            }
+        }
+
+        Ok(())
+    }
 }
 
 pub fn validate_static_schema(

src/prism/logstream/mod.rs

@@ -28,7 +28,6 @@ use crate::{
     handlers::http::{
         cluster::utils::{merge_quried_stats, IngestionStats, QueriedStats, StorageStats},
         logstream::error::StreamError,
-        query::update_schema_when_distributed,
     },
     parseable::{StreamNotFound, PARSEABLE},
     stats,
@@ -77,7 +76,10 @@ async fn get_stream_schema_helper(stream_name: &str) -> Result<Arc<Schema>, Stre
     }
 
     let stream = PARSEABLE.get_stream(stream_name)?;
-    match update_schema_when_distributed(&vec![stream_name.to_owned()]).await {
+    match PARSEABLE
+        .update_schema_when_distributed(&vec![stream_name.to_owned()])
+        .await
+    {
         Ok(_) => {
             let schema = stream.get_schema();
             Ok(schema)