refactor: share Stream state when processing

Author: Devdutt Shenoi

src/event/format/json.rs

@@ -43,7 +43,7 @@ use crate::{
     kinesis::{flatten_kinesis_logs, Message},
     metadata::SchemaVersion,
     otel::{logs::flatten_otel_logs, metrics::flatten_otel_metrics, traces::flatten_otel_traces},
-    storage::StreamType,
+    parseable::Stream,
     utils::{
         arrow::get_field,
         json::{flatten_json_body, Json},
@@ -224,23 +224,24 @@ impl EventFormat for Event {
     /// Converts a JSON event into a Parseable Event
     fn into_event(
         self,
-        stream_name: String,
         origin_size: u64,
-        storage_schema: &HashMap<String, Arc<Field>>,
-        static_schema_flag: bool,
-        custom_partitions: Option<&String>,
-        time_partition: Option<&String>,
-        time_partition_limit: Option<NonZeroU32>,
-        schema_version: SchemaVersion,
+        stream: &Stream,
         log_source: &LogSource,
-        stream_type: StreamType,
     ) -> anyhow::Result<super::Event> {
+        let time_partition = stream.get_time_partition();
+        let time_partition_limit = stream.get_time_partition_limit();
+        let static_schema_flag = stream.get_static_schema_flag();
+        let custom_partitions = stream.get_custom_partition();
+        let schema_version = stream.get_schema_version();
+        let storage_schema = stream.get_schema_raw();
+        let stream_type = stream.get_stream_type();
+
         let p_timestamp = self.p_timestamp;
         let (data, schema, is_first_event) = self.to_data(
-            storage_schema,
-            time_partition,
+            &storage_schema,
+            time_partition.as_ref(),
             time_partition_limit,
-            custom_partitions,
+            custom_partitions.as_ref(),
             schema_version,
             log_source,
         )?;
@@ -255,18 +256,18 @@ impl EventFormat for Event {
                 None => HashMap::new(),
             };
 
-            let parsed_timestamp = match time_partition {
-                Some(time_partition) => extract_and_parse_time(&json, time_partition)?,
+            let parsed_timestamp = match time_partition.as_ref() {
+                Some(time_partition) => extract_and_parse_time(&json, time_partition.as_ref())?,
                 _ => p_timestamp.naive_utc(),
             };
 
             let rb = Self::into_recordbatch(
                 p_timestamp,
                 vec![json],
                 schema.clone(),
-                storage_schema,
+                &storage_schema,
                 static_schema_flag,
-                time_partition,
+                time_partition.as_ref(),
                 schema_version,
             )?;
 
@@ -278,7 +279,6 @@ impl EventFormat for Event {
         }
 
         Ok(super::Event {
-            stream_name,
             origin_format: "json",
             origin_size,
             is_first_event,

src/event/format/mod.rs

@@ -33,7 +33,7 @@ use serde_json::Value;
 
 use crate::{
     metadata::SchemaVersion,
-    storage::StreamType,
+    parseable::Stream,
     utils::{
         arrow::{get_field, get_timestamp_array, replace_columns},
         json::Json,
@@ -165,19 +165,11 @@ pub trait EventFormat: Sized {
         Ok(rb)
     }
 
-    #[allow(clippy::too_many_arguments)]
     fn into_event(
         self,
-        stream_name: String,
         origin_size: u64,
-        storage_schema: &HashMap<String, Arc<Field>>,
-        static_schema_flag: bool,
-        custom_partitions: Option<&String>,
-        time_partition: Option<&String>,
-        time_partition_limit: Option<NonZeroU32>,
-        schema_version: SchemaVersion,
+        stream: &Stream,
         log_source: &LogSource,
-        stream_type: StreamType,
     ) -> Result<Event, AnyError>;
 }
 

src/event/mod.rs

@@ -27,7 +27,7 @@ use std::sync::Arc;
 use self::error::EventError;
 use crate::{
     metadata::update_stats,
-    parseable::{StagingError, PARSEABLE},
+    parseable::{StagingError, Stream, PARSEABLE},
     storage::StreamType,
     LOCK_EXPECT,
 };
@@ -43,7 +43,6 @@ pub struct PartitionEvent {
 }
 
 pub struct Event {
-    pub stream_name: String,
     pub origin_format: &'static str,
     pub origin_size: u64,
     pub is_first_event: bool,
@@ -54,7 +53,7 @@ pub struct Event {
 
 // Events holds the schema related to a each event for a single log stream
 impl Event {
-    pub fn process(self) -> Result<(), EventError> {
+    pub fn process(self, stream: &Stream) -> Result<(), EventError> {
         for partition in self.partitions {
             let mut key = get_schema_key(&partition.rb.schema().fields);
             if self.time_partition.is_some() {
@@ -72,10 +71,10 @@ impl Event {
             }
 
             if self.is_first_event {
-                commit_schema(&self.stream_name, partition.rb.schema())?;
+                commit_schema(&stream.stream_name, partition.rb.schema())?;
             }
 
-            PARSEABLE.get_or_create_stream(&self.stream_name).push(
+            stream.push(
                 &key,
                 &partition.rb,
                 partition.parsed_timestamp,
@@ -84,23 +83,23 @@ impl Event {
             )?;
 
             update_stats(
-                &self.stream_name,
+                &stream.stream_name,
                 self.origin_format,
                 self.origin_size,
                 partition.rb.num_rows(),
                 partition.parsed_timestamp.date(),
             );
 
-            crate::livetail::LIVETAIL.process(&self.stream_name, &partition.rb);
+            crate::livetail::LIVETAIL.process(&stream.stream_name, &partition.rb);
         }
         Ok(())
     }
 
-    pub fn process_unchecked(&self) -> Result<(), EventError> {
+    pub fn process_unchecked(&self, stream: &Stream) -> Result<(), EventError> {
         for partition in &self.partitions {
             let key = get_schema_key(&partition.rb.schema().fields);
 
-            PARSEABLE.get_or_create_stream(&self.stream_name).push(
+            stream.push(
                 &key,
                 &partition.rb,
                 partition.parsed_timestamp,

src/handlers/http/ingest.rs

@@ -31,7 +31,7 @@ use crate::event::format::LogSource;
 use crate::event::{self, PartitionEvent};
 use crate::handlers::{LOG_SOURCE_KEY, STREAM_NAME_HEADER_KEY};
 use crate::option::Mode;
-use crate::parseable::{StreamNotFound, PARSEABLE};
+use crate::parseable::{Stream, StreamNotFound, PARSEABLE};
 use crate::storage::{ObjectStorageError, StreamType};
 use crate::utils::header_parsing::ParseHeaderError;
 use crate::utils::json::flatten::JsonFlattenError;
@@ -240,10 +240,9 @@ pub async fn post_event(
 
 pub async fn push_logs_unchecked(
     batches: RecordBatch,
-    stream_name: &str,
+    stream: &Stream,
 ) -> Result<event::Event, PostError> {
     let unchecked_event = event::Event {
-        stream_name: stream_name.to_string(),
         origin_format: "json",
         origin_size: 0,
         time_partition: None,
@@ -255,7 +254,7 @@ pub async fn push_logs_unchecked(
         }],
         stream_type: StreamType::UserDefined,
     };
-    unchecked_event.process_unchecked()?;
+    unchecked_event.process_unchecked(stream)?;
 
     Ok(unchecked_event)
 }

src/parseable/mod.rs

@@ -28,7 +28,7 @@ use http::{header::CONTENT_TYPE, HeaderName, HeaderValue, StatusCode};
 use once_cell::sync::Lazy;
 pub use staging::StagingError;
 use streams::StreamRef;
-pub use streams::{StreamNotFound, Streams};
+pub use streams::{Stream, StreamNotFound, Streams};
 use tracing::error;
 
 #[cfg(feature = "kafka")]

src/parseable/streams.rs

@@ -114,30 +114,11 @@ impl Stream {
     }
 
     pub async fn push_logs(&self, json: Value, log_source: &LogSource) -> anyhow::Result<()> {
-        let time_partition = self.get_time_partition();
-        let time_partition_limit = self.get_time_partition_limit();
-        let static_schema_flag = self.get_static_schema_flag();
-        let custom_partition = self.get_custom_partition();
-        let schema_version = self.get_schema_version();
-        let schema = self.get_schema_raw();
-        let stream_type = self.get_stream_type();
-
         let origin_size = serde_json::to_vec(&json).unwrap().len() as u64; // string length need not be the same as byte length
 
         json::Event::new(json)
-            .into_event(
-                self.stream_name.to_owned(),
-                origin_size,
-                &schema,
-                static_schema_flag,
-                custom_partition.as_ref(),
-                time_partition.as_ref(),
-                time_partition_limit,
-                schema_version,
-                log_source,
-                stream_type,
-            )?
-            .process()?;
+            .into_event(origin_size, self, log_source)?
+            .process(self)?;
 
         Ok(())
     }

src/utils/arrow/flight.rs

@@ -95,14 +95,14 @@ pub async fn append_temporary_events(
     Event,
     Status,
 > {
-    let schema = PARSEABLE
+    let stream = PARSEABLE
         .get_stream(stream_name)
-        .map_err(|err| Status::failed_precondition(format!("Metadata Error: {}", err)))?
-        .get_schema();
+        .map_err(|err| Status::failed_precondition(format!("Metadata Error: {}", err)))?;
+    let schema = stream.get_schema();
     let rb = concat_batches(&schema, minute_result)
         .map_err(|err| Status::failed_precondition(format!("ArrowError: {}", err)))?;
 
-    let event = push_logs_unchecked(rb, stream_name)
+    let event = push_logs_unchecked(rb, &stream)
         .await
         .map_err(|err| Status::internal(err.to_string()))?;
     Ok(event)