test + refactor: staging (#1129)

Author: Devdutt Shenoi

.gitignore

@@ -1,6 +1,6 @@
 target
 data*
-staging/
+staging/*
 limitcache
 examples
 cert.pem

Cargo.lock

@@ -129,6 +129,7 @@ zip = { version = "2.2.0", default-features = false, features = ["deflate"] }
 [dev-dependencies]
 rstest = "0.23.0"
 arrow = "53.0.0"
+temp-dir = "0.1.14"
 
 [package.metadata.parseable_ui]
 assets-url = "https://github.com/parseablehq/console/releases/download/v0.9.18/build.zip"

Cargo.toml

@@ -102,7 +102,7 @@ pub struct BlobStoreArgs {
     pub storage: AzureBlobConfig,
 }
 
-#[derive(Parser, Debug)]
+#[derive(Parser, Debug, Default)]
 pub struct Options {
     // Authentication
     #[arg(long, env = "P_USERNAME", help = "Admin username to be set for this Parseable server", default_value = DEFAULT_USERNAME)]
@@ -283,7 +283,7 @@ pub struct Options {
     pub ingestor_endpoint: String,
 
     #[command(flatten)]
-    oidc: Option<OidcConfig>,
+    pub oidc: Option<OidcConfig>,
 
     // Kafka configuration (conditionally compiled)
     #[cfg(any(

src/cli.rs

@@ -18,16 +18,14 @@
 */
 
 pub mod format;
-mod writer;
 
 use arrow_array::RecordBatch;
 use arrow_schema::{Field, Fields, Schema};
 use itertools::Itertools;
 use std::sync::Arc;
 
 use self::error::EventError;
-pub use self::writer::STREAM_WRITERS;
-use crate::{metadata, storage::StreamType};
+use crate::{metadata, staging::STAGING, storage::StreamType};
 use chrono::NaiveDateTime;
 use std::collections::HashMap;
 
@@ -52,36 +50,32 @@ impl Event {
         let mut key = get_schema_key(&self.rb.schema().fields);
         if self.time_partition.is_some() {
             let parsed_timestamp_to_min = self.parsed_timestamp.format("%Y%m%dT%H%M").to_string();
-            key = format!("{key}{parsed_timestamp_to_min}");
+            key.push_str(&parsed_timestamp_to_min);
         }
 
         if !self.custom_partition_values.is_empty() {
-            let mut custom_partition_key = String::default();
             for (k, v) in self.custom_partition_values.iter().sorted_by_key(|v| v.0) {
-                custom_partition_key = format!("{custom_partition_key}&{k}={v}");
+                key.push_str(&format!("&{k}={v}"));
             }
-            key = format!("{key}{custom_partition_key}");
         }
 
-        let num_rows = self.rb.num_rows() as u64;
         if self.is_first_event {
             commit_schema(&self.stream_name, self.rb.schema())?;
         }
 
-        STREAM_WRITERS.append_to_local(
-            &self.stream_name,
+        STAGING.get_or_create_stream(&self.stream_name).push(
             &key,
             &self.rb,
             self.parsed_timestamp,
             &self.custom_partition_values,
-            &self.stream_type,
+            self.stream_type,
         )?;
 
         metadata::STREAM_INFO.update_stats(
             &self.stream_name,
             self.origin_format,
             self.origin_size,
-            num_rows,
+            self.rb.num_rows(),
             self.parsed_timestamp,
         )?;
 
@@ -93,21 +87,16 @@ impl Event {
     pub fn process_unchecked(&self) -> Result<(), EventError> {
         let key = get_schema_key(&self.rb.schema().fields);
 
-        STREAM_WRITERS.append_to_local(
-            &self.stream_name,
+        STAGING.get_or_create_stream(&self.stream_name).push(
             &key,
             &self.rb,
             self.parsed_timestamp,
             &self.custom_partition_values,
-            &self.stream_type,
+            self.stream_type,
         )?;
 
         Ok(())
     }
-
-    pub fn clear(&self, stream_name: &str) {
-        STREAM_WRITERS.clear(stream_name);
-    }
 }
 
 pub fn get_schema_key(fields: &[Arc<Field>]) -> String {
@@ -138,14 +127,13 @@ pub mod error {
     use arrow_schema::ArrowError;
 
     use crate::metadata::error::stream_info::MetadataError;
+    use crate::staging::StagingError;
     use crate::storage::ObjectStorageError;
 
-    use super::writer::errors::StreamWriterError;
-
     #[derive(Debug, thiserror::Error)]
     pub enum EventError {
         #[error("Stream Writer Failed: {0}")]
-        StreamWriter(#[from] StreamWriterError),
+        StreamWriter(#[from] StagingError),
         #[error("Metadata Error: {0}")]
         Metadata(#[from] MetadataError),
         #[error("Stream Writer Failed: {0}")]