Merge branch 'main' into disk-writer

Author: Devdutt Shenoi

src/connectors/kafka/processor.rs

@@ -28,7 +28,7 @@ use tracing::{debug, error};
 use crate::{
     connectors::common::processor::Processor,
     event::{
-        format::{json, EventFormat, LogSource},
+        format::{json, EventFormat, LogSourceEntry},
         Event as ParseableEvent,
     },
     parseable::PARSEABLE,
@@ -49,9 +49,14 @@ impl ParseableSinkProcessor {
             .first()
             .map(|r| r.topic.as_str())
             .unwrap_or_default();
+        let log_source_entry = LogSourceEntry::default();
 
         PARSEABLE
-            .create_stream_if_not_exists(stream_name, StreamType::UserDefined, LogSource::Json)
+            .create_stream_if_not_exists(
+                stream_name,
+                StreamType::UserDefined,
+                vec![log_source_entry],
+            )
             .await?;
 
         let stream = PARSEABLE.get_stream(stream_name)?;

src/event/format/json.rs

@@ -23,7 +23,7 @@ use anyhow::anyhow;
 use arrow_array::RecordBatch;
 use arrow_json::reader::{infer_json_schema_from_iterator, ReaderBuilder};
 use arrow_schema::{DataType, Field, Fields, Schema};
-use chrono::{DateTime, NaiveDateTime, Utc};
+use chrono::{DateTime, NaiveDate, NaiveDateTime, Utc};
 use datafusion::arrow::util::bit_util::round_upto_multiple_of_64;
 use itertools::Itertools;
 use serde_json::Value;
@@ -62,6 +62,7 @@ impl EventFormat for Event {
         schema: &HashMap<String, Arc<Field>>,
         time_partition: Option<&String>,
         schema_version: SchemaVersion,
+        static_schema_flag: bool,
     ) -> Result<(Self::Data, Vec<Arc<Field>>, bool), anyhow::Error> {
         let stream_schema = schema;
 
@@ -111,7 +112,7 @@ impl EventFormat for Event {
 
         if value_arr
             .iter()
-            .any(|value| fields_mismatch(&schema, value, schema_version))
+            .any(|value| fields_mismatch(&schema, value, schema_version, static_schema_flag))
         {
             return Err(anyhow!(
                 "Could not process this event due to mismatch in datatype"
@@ -253,73 +254,131 @@ fn collect_keys<'a>(values: impl Iterator<Item = &'a Value>) -> Result<Vec<&'a s
     Ok(keys)
 }
 
-fn fields_mismatch(schema: &[Arc<Field>], body: &Value, schema_version: SchemaVersion) -> bool {
+fn fields_mismatch(
+    schema: &[Arc<Field>],
+    body: &Value,
+    schema_version: SchemaVersion,
+    static_schema_flag: bool,
+) -> bool {
     for (name, val) in body.as_object().expect("body is of object variant") {
         if val.is_null() {
             continue;
         }
         let Some(field) = get_field(schema, name) else {
             return true;
         };
-        if !valid_type(field.data_type(), val, schema_version) {
+        if !valid_type(field, val, schema_version, static_schema_flag) {
             return true;
         }
     }
     false
 }
 
-fn valid_type(data_type: &DataType, value: &Value, schema_version: SchemaVersion) -> bool {
-    match data_type {
+fn valid_type(
+    field: &Field,
+    value: &Value,
+    schema_version: SchemaVersion,
+    static_schema_flag: bool,
+) -> bool {
+    match field.data_type() {
         DataType::Boolean => value.is_boolean(),
-        DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => value.is_i64(),
+        DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => {
+            validate_int(value, static_schema_flag)
+        }
         DataType::UInt8 | DataType::UInt16 | DataType::UInt32 | DataType::UInt64 => value.is_u64(),
         DataType::Float16 | DataType::Float32 => value.is_f64(),
-        // All numbers can be cast as Float64 from schema version v1
-        DataType::Float64 if schema_version == SchemaVersion::V1 => value.is_number(),
-        DataType::Float64 if schema_version != SchemaVersion::V1 => value.is_f64(),
+        DataType::Float64 => validate_float(value, schema_version, static_schema_flag),
         DataType::Utf8 => value.is_string(),
-        DataType::List(field) => {
-            let data_type = field.data_type();
-            if let Value::Array(arr) = value {
-                for elem in arr {
-                    if elem.is_null() {
-                        continue;
-                    }
-                    if !valid_type(data_type, elem, schema_version) {
-                        return false;
-                    }
-                }
-            }
-            true
-        }
+        DataType::List(field) => validate_list(field, value, schema_version, static_schema_flag),
         DataType::Struct(fields) => {
-            if let Value::Object(val) = value {
-                for (key, value) in val {
-                    let field = (0..fields.len())
-                        .find(|idx| fields[*idx].name() == key)
-                        .map(|idx| &fields[idx]);
-
-                    if let Some(field) = field {
-                        if value.is_null() {
-                            continue;
-                        }
-                        if !valid_type(field.data_type(), value, schema_version) {
-                            return false;
-                        }
-                    } else {
-                        return false;
-                    }
-                }
-                true
-            } else {
-                false
+            validate_struct(fields, value, schema_version, static_schema_flag)
+        }
+        DataType::Date32 => {
+            if let Value::String(s) = value {
+                return NaiveDate::parse_from_str(s, "%Y-%m-%d").is_ok();
             }
+            false
         }
         DataType::Timestamp(_, _) => value.is_string() || value.is_number(),
         _ => {
-            error!("Unsupported datatype {:?}, value {:?}", data_type, value);
-            unreachable!()
+            error!(
+                "Unsupported datatype {:?}, value {:?}",
+                field.data_type(),
+                value
+            );
+            false
+        }
+    }
+}
+
+fn validate_int(value: &Value, static_schema_flag: bool) -> bool {
+    // allow casting string to int for static schema
+    if static_schema_flag {
+        if let Value::String(s) = value {
+            return s.trim().parse::<i64>().is_ok();
+        }
+    }
+    value.is_i64()
+}
+
+fn validate_float(value: &Value, schema_version: SchemaVersion, static_schema_flag: bool) -> bool {
+    // allow casting string to int for static schema
+    if static_schema_flag {
+        if let Value::String(s) = value.clone() {
+            let trimmed = s.trim();
+            return trimmed.parse::<f64>().is_ok() || trimmed.parse::<i64>().is_ok();
+        }
+        return value.is_number();
+    }
+    match schema_version {
+        SchemaVersion::V1 => value.is_number(),
+        _ => value.is_f64(),
+    }
+}
+
+fn validate_list(
+    field: &Field,
+    value: &Value,
+    schema_version: SchemaVersion,
+    static_schema_flag: bool,
+) -> bool {
+    if let Value::Array(arr) = value {
+        for elem in arr {
+            if elem.is_null() {
+                continue;
+            }
+            if !valid_type(field, elem, schema_version, static_schema_flag) {
+                return false;
+            }
+        }
+    }
+    true
+}
+
+fn validate_struct(
+    fields: &Fields,
+    value: &Value,
+    schema_version: SchemaVersion,
+    static_schema_flag: bool,
+) -> bool {
+    if let Value::Object(val) = value {
+        for (key, value) in val {
+            let field = fields.iter().find(|f| f.name() == key);
+
+            if let Some(field) = field {
+                if value.is_null() {
+                    continue;
+                }
+                if !valid_type(field, value, schema_version, static_schema_flag) {
+                    return false;
+                }
+            } else {
+                return false;
+            }
         }
+        true
+    } else {
+        false
     }
 }
 

src/event/format/mod.rs

@@ -44,7 +44,7 @@ static TIME_FIELD_NAME_PARTS: [&str; 2] = ["time", "date"];
 type EventSchema = Vec<Arc<Field>>;
 
 /// Source of the logs, used to perform special processing for certain sources
-#[derive(Default, Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Default, Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Hash)]
 pub enum LogSource {
     // AWS Kinesis sends logs in the format of a json array
     Kinesis,
@@ -92,6 +92,23 @@ impl Display for LogSource {
     }
 }
 
+/// Contains the format name and a list of known field names that are associated with the said format.
+/// Stored on disk as part of `ObjectStoreFormat` in stream.json
+#[derive(Default, Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct LogSourceEntry {
+    pub log_source_format: LogSource,
+    pub fields: HashSet<String>,
+}
+
+impl LogSourceEntry {
+    pub fn new(log_source_format: LogSource, fields: HashSet<String>) -> Self {
+        LogSourceEntry {
+            log_source_format,
+            fields,
+        }
+    }
+}
+
 // Global Trait for event format
 // This trait is implemented by all the event formats
 pub trait EventFormat: Sized {
@@ -102,6 +119,7 @@ pub trait EventFormat: Sized {
         schema: &HashMap<String, Arc<Field>>,
         time_partition: Option<&String>,
         schema_version: SchemaVersion,
+        static_schema_flag: bool,
     ) -> Result<(Self::Data, EventSchema, bool), AnyError>;
 
     fn decode(data: Self::Data, schema: Arc<Schema>) -> Result<RecordBatch, AnyError>;
@@ -117,8 +135,12 @@ pub trait EventFormat: Sized {
         schema_version: SchemaVersion,
     ) -> Result<(RecordBatch, bool), AnyError> {
         let p_timestamp = self.get_p_timestamp();
-        let (data, mut schema, is_first) =
-            self.to_data(storage_schema, time_partition, schema_version)?;
+        let (data, mut schema, is_first) = self.to_data(
+            storage_schema,
+            time_partition,
+            schema_version,
+            static_schema_flag,
+        )?;
 
         if get_field(&schema, DEFAULT_TIMESTAMP_KEY).is_some() {
             return Err(anyhow!(

src/handlers/http/ingest.rs

@@ -16,7 +16,7 @@
  *
  */
 
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 
 use actix_web::web::{Json, Path};
 use actix_web::{http::header::ContentType, HttpRequest, HttpResponse};
@@ -28,10 +28,13 @@ use serde_json::Value;
 
 use crate::event;
 use crate::event::error::EventError;
-use crate::event::format::{self, EventFormat, LogSource};
+use crate::event::format::{self, EventFormat, LogSource, LogSourceEntry};
 use crate::handlers::{LOG_SOURCE_KEY, STREAM_NAME_HEADER_KEY};
 use crate::metadata::SchemaVersion;
 use crate::option::Mode;
+use crate::otel::logs::OTEL_LOG_KNOWN_FIELD_LIST;
+use crate::otel::metrics::OTEL_METRICS_KNOWN_FIELD_LIST;
+use crate::otel::traces::OTEL_TRACES_KNOWN_FIELD_LIST;
 use crate::parseable::{StreamNotFound, PARSEABLE};
 use crate::storage::{ObjectStorageError, StreamType};
 use crate::utils::header_parsing::ParseHeaderError;
@@ -55,9 +58,6 @@ pub async fn ingest(req: HttpRequest, Json(json): Json<Value>) -> Result<HttpRes
     if internal_stream_names.contains(&stream_name) {
         return Err(PostError::InternalStream(stream_name));
     }
-    PARSEABLE
-        .create_stream_if_not_exists(&stream_name, StreamType::UserDefined, LogSource::default())
-        .await?;
 
     let log_source = req
         .headers()
@@ -72,6 +72,15 @@ pub async fn ingest(req: HttpRequest, Json(json): Json<Value>) -> Result<HttpRes
         return Err(PostError::OtelNotSupported);
     }
 
+    let log_source_entry = LogSourceEntry::new(log_source.clone(), HashSet::new());
+    PARSEABLE
+        .create_stream_if_not_exists(
+            &stream_name,
+            StreamType::UserDefined,
+            vec![log_source_entry],
+        )
+        .await?;
+
     flatten_and_push_logs(json, &stream_name, &log_source).await?;
 
     Ok(HttpResponse::Ok().finish())
@@ -119,8 +128,20 @@ pub async fn handle_otel_logs_ingestion(
     }
 
     let stream_name = stream_name.to_str().unwrap().to_owned();
+
+    let log_source_entry = LogSourceEntry::new(
+        log_source.clone(),
+        OTEL_LOG_KNOWN_FIELD_LIST
+            .iter()
+            .map(|&s| s.to_string())
+            .collect(),
+    );
     PARSEABLE
-        .create_stream_if_not_exists(&stream_name, StreamType::UserDefined, LogSource::OtelLogs)
+        .create_stream_if_not_exists(
+            &stream_name,
+            StreamType::UserDefined,
+            vec![log_source_entry],
+        )
         .await?;
 
     flatten_and_push_logs(json, &stream_name, &log_source).await?;
@@ -146,11 +167,18 @@ pub async fn handle_otel_metrics_ingestion(
         return Err(PostError::IncorrectLogSource(LogSource::OtelMetrics));
     }
     let stream_name = stream_name.to_str().unwrap().to_owned();
+    let log_source_entry = LogSourceEntry::new(
+        log_source.clone(),
+        OTEL_METRICS_KNOWN_FIELD_LIST
+            .iter()
+            .map(|&s| s.to_string())
+            .collect(),
+    );
     PARSEABLE
         .create_stream_if_not_exists(
             &stream_name,
             StreamType::UserDefined,
-            LogSource::OtelMetrics,
+            vec![log_source_entry],
         )
         .await?;
 
@@ -178,8 +206,20 @@ pub async fn handle_otel_traces_ingestion(
         return Err(PostError::IncorrectLogSource(LogSource::OtelTraces));
     }
     let stream_name = stream_name.to_str().unwrap().to_owned();
+    let log_source_entry = LogSourceEntry::new(
+        log_source.clone(),
+        OTEL_TRACES_KNOWN_FIELD_LIST
+            .iter()
+            .map(|&s| s.to_string())
+            .collect(),
+    );
+
     PARSEABLE
-        .create_stream_if_not_exists(&stream_name, StreamType::UserDefined, LogSource::OtelTraces)
+        .create_stream_if_not_exists(
+            &stream_name,
+            StreamType::UserDefined,
+            vec![log_source_entry],
+        )
         .await?;
 
     flatten_and_push_logs(json, &stream_name, &log_source).await?;