Add support for defining schemas to be used by the sinks

Right now the Kafka sink does not support the use of the defined
schemas, but these allow for defining valid/acceptable schemas up front
for data written to specific topics.

What this will _not_ do however is any form of type coercion! Make sure
the schemas are the right types for the data coming in!

Author: rtyler

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "hotdog"
-version = "1.0.2"
+version = "1.1.0"
 authors = ["R. Tyler Croy <rtyler@buoyantdata.com>"]
 edition = "2024"
 
@@ -47,6 +47,7 @@ rdkafka = { version = "0", features = ["ssl", "sasl"]}
 object_store = { version = "0.12", features = ["aws", "cloud"] }
 arrow-array = "55"
 arrow-json = "55"
+arrow-schema = "55"
 parquet = { version = "55", features = ["arrow", "async", "object_store"]}
 
 # Used for rule matching

src/connection.rs

@@ -84,7 +84,7 @@ impl Connection {
 
         while let Some(line) = lines.next().await {
             let line = line?;
-            debug!("log: {}", line);
+            trace!("log: {}", line);
 
             let parsed = parse::parse_line(line);
 
@@ -102,7 +102,7 @@ impl Connection {
             let mut msg = parsed.unwrap();
             self.stats.counter(Stats::LineReceived.into()).count(1);
             let mut continue_rules = true;
-            debug!("parsed as: {}", msg.msg);
+            trace!("parsed as: {}", msg.msg);
 
             for rule in self.settings.rules.iter() {
                 /*

src/main.rs

@@ -18,6 +18,7 @@ mod json;
 mod merge;
 mod parse;
 mod rules;
+mod schema;
 mod serve;
 mod settings;
 mod sink;

src/schema.rs

@@ -0,0 +1,196 @@
+//!
+//! Simple schema mappings to help turn defined schemas into Arrow schemas
+
+use arrow_schema::DataType;
+use serde::Deserialize;
+use tracing::log::*;
+
+use std::collections::HashMap;
+
+/// Supported schema definition types for the configuration of hotdog
+#[derive(Clone, Debug, Deserialize, PartialEq)]
+#[serde(rename_all = "lowercase")]
+pub enum FieldType {
+    String,
+    Struct,
+    Long,
+    Integer,
+    Timestamp,
+    Float,
+    Boolean,
+}
+
+/// Helpful converter for [arrow_schema::Schema] conversion
+impl Into<DataType> for &FieldType {
+    fn into(self) -> DataType {
+        match self {
+            FieldType::String => DataType::Utf8,
+            FieldType::Boolean => DataType::Boolean,
+            FieldType::Integer => DataType::Int32,
+            FieldType::Long => DataType::Int64,
+            FieldType::Float => DataType::Float64,
+            FieldType::Timestamp => DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, None),
+            unknown => {
+                warn!(
+                    "Asked to convert a schema type which cannot be automatically converted: {unknown:?}"
+                );
+                DataType::Null
+            }
+        }
+    }
+}
+
+/// A field in the schema, which supports recursion. The name of the field is expected to be known
+/// by the container of the field
+#[derive(Clone, Debug, Deserialize)]
+pub struct Field {
+    r#type: FieldType,
+    fields: Option<HashMap<String, Field>>,
+}
+
+impl Field {
+    /// Retrieve the given [FieldType]
+    pub fn get_type(&self) -> &FieldType {
+        &self.r#type
+    }
+
+    /// Retrieve the nested fields should they exist
+    pub fn fields(&self) -> Option<&HashMap<String, Field>> {
+        self.fields.as_ref()
+    }
+}
+
+/// Convert the given set of fields to an [arrow_schema::Schema] for use elsewher einside of hotdog
+pub fn into_arrow_schema(
+    fields: &HashMap<String, Field>,
+) -> Result<arrow_schema::Schema, arrow_schema::ArrowError> {
+    let mut arrow_fields = vec![];
+
+    for (name, field) in fields {
+        if field.get_type() == &FieldType::Struct {
+            if let Some(fields) = field.fields.as_ref() {
+                let struct_fields = into_arrow_schema(fields)?;
+                arrow_fields.push(arrow_schema::Field::new(
+                    name,
+                    DataType::Struct(struct_fields.fields),
+                    true,
+                ));
+            } else {
+                error!(
+                    "I don't know how to handle structs without fields! Calling it null :shrug:"
+                );
+                arrow_fields.push(arrow_schema::Field::new(name, DataType::Null, true));
+            }
+        } else {
+            arrow_fields.push(arrow_schema::Field::new(
+                name,
+                field.get_type().into(),
+                true,
+            ));
+        }
+    }
+
+    Ok(arrow_schema::Schema::new(arrow_fields))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_deser_simple_field() {
+        let buf = r#"
+---
+version:
+  type: string
+"#;
+        let fields: HashMap<String, Field> =
+            serde_yaml::from_str(buf).expect("Failed to deserialize a simple string field");
+
+        if let Some(field) = fields.get("version") {
+            assert_eq!(field.get_type(), &FieldType::String);
+        } else {
+            panic!("Failed to get a version field");
+        }
+    }
+
+    #[test]
+    fn test_deser_nested_map() {
+        let buf = r#"
+---
+meta:
+  type: struct
+  fields:
+    version:
+      type: string
+"#;
+        let fields: HashMap<String, Field> =
+            serde_yaml::from_str(buf).expect("Failed to deserialize a nested field");
+
+        if let Some(field) = fields.get("meta") {
+            assert_eq!(field.get_type(), &FieldType::Struct);
+
+            let nested = field
+                .fields
+                .as_ref()
+                .expect("Failed to get the nested fields");
+            assert_eq!(
+                nested.get("version").unwrap().get_type(),
+                &FieldType::String
+            );
+        } else {
+            panic!("Failed to get a timestamp field");
+        }
+    }
+
+    /// The whole point of this schema definition is to get to an arrow schema definition!
+    #[test]
+    fn test_convert_simple_to_arrow() {
+        let buf = r#"
+---
+version:
+    type: string
+"#;
+        let fields: HashMap<String, Field> =
+            serde_yaml::from_str(buf).expect("Failed to deserialize a nested field");
+
+        let arr: arrow_schema::Schema =
+            into_arrow_schema(&fields).expect("This is a valid arrow schema");
+        assert_ne!(
+            arr.fields.len(),
+            0,
+            "The converted schema should have one field"
+        );
+    }
+
+    #[test]
+    fn test_convert_nested_to_arrow() {
+        let buf = r#"
+---
+meta:
+  type: struct
+  fields:
+    version:
+        type: string
+"#;
+        let fields: HashMap<String, Field> =
+            serde_yaml::from_str(buf).expect("Failed to deserialize a nested field");
+
+        let arr: arrow_schema::Schema =
+            into_arrow_schema(&fields).expect("This is a valid arrow schema");
+
+        let (_size, field) = arr
+            .fields
+            .find("meta")
+            .expect("Failed to find the meta field");
+
+        let struct_fields = arrow_schema::Fields::from(vec![arrow_schema::Field::new(
+            "version",
+            DataType::Utf8,
+            true,
+        )]);
+        let struct_type = DataType::Struct(struct_fields);
+
+        assert_eq!(field.data_type(), &struct_type);
+    }
+}

src/serve/mod.rs

@@ -92,7 +92,11 @@ pub trait Server {
         // If the Kafka sink is defined in the configuration, then spin up the configuration
         if let Some(kafka_conf) = &state.settings.global.kafka {
             info!("Configuring a Kafka sink with: {kafka_conf:?}");
-            let mut kafka = Kafka::new(kafka_conf.clone(), state.stats.clone());
+            let mut kafka = Kafka::new(
+                kafka_conf.clone(),
+                &state.settings.schemas,
+                state.stats.clone(),
+            );
 
             kafka.bootstrap().await;
             sender = Some(kafka.get_sender());
@@ -106,7 +110,11 @@ pub trait Server {
 
         if let Some(parquet_conf) = &state.settings.global.parquet {
             info!("Configuring a Parquet sink with: {parquet_conf:?}");
-            let pq = Parquet::new(parquet_conf.clone(), state.stats.clone());
+            let pq = Parquet::new(
+                parquet_conf.clone(),
+                &state.settings.schemas,
+                state.stats.clone(),
+            );
             sender = Some(pq.get_sender());
             smol::spawn(async move {
                 debug!("Starting Parquet loop");

src/settings.rs

@@ -4,6 +4,7 @@ use serde_json::Value;
 use tracing::log::*;
 use uuid::Uuid;
 
+use std::collections::HashMap;
 use std::path::Path;
 
 pub fn load(file: &str) -> Settings {
@@ -159,23 +160,37 @@ pub struct Global {
     pub status: Option<Status>,
 }
 
+/// Schema is a definition of an Arrow schema which can be used for deserialization
+#[derive(Debug, Deserialize)]
+pub struct Schema {
+    /// Output topic this schema applies to
+    pub topic: String,
+    /// Fields of the actual schema
+    pub fields: HashMap<String, crate::schema::Field>,
+}
+
 #[derive(Debug, Deserialize)]
 pub struct Settings {
     pub global: Global,
     pub rules: Vec<Rule>,
+    #[serde(default = "default_schemas")]
+    /// Optionally defined schemas for this configuration
+    pub schemas: Vec<Schema>,
 }
 
 impl Settings {
-    /**
-     * Populate any configuration caches which we want to us
-     */
+    /// Populate any configuration caches which we want to us
     fn populate_caches(&mut self) {
         self.rules.iter_mut().for_each(|rule| {
             rule.populate_caches();
         });
     }
 }
 
+fn default_schemas() -> Vec<Schema> {
+    vec![]
+}
+
 fn default_none<T>() -> Option<T> {
     None
 }

src/sink/kafka.rs

@@ -41,7 +41,7 @@ pub struct Kafka {
 impl Sink for Kafka {
     type Config = Config;
 
-    fn new(config: Config, stats: InputQueueScope) -> Self {
+    fn new(config: Config, _schemas: &[crate::settings::Schema], stats: InputQueueScope) -> Self {
         let (tx, rx) = bounded(config.buffer);
         Kafka {
             producer: None,
@@ -242,7 +242,7 @@ mod tests {
 
         let bucket = dipstick::AtomicBucket::new();
         let stats = InputQueueScope::wrap(bucket.clone(), 100);
-        let mut k = Kafka::new(config, stats);
+        let mut k = Kafka::new(config, &[], stats);
         smol::block_on(k.bootstrap());
     }
 

src/sink/mod.rs

@@ -24,7 +24,11 @@ pub trait Sink: Send + Sync {
     /// Construct the Sink.
     ///
     /// This function should not do anything but initialize settings and variables
-    fn new(config: Self::Config, stats: InputQueueScope) -> Self;
+    fn new(
+        config: Self::Config,
+        schemas: &[crate::settings::Schema],
+        stats: InputQueueScope,
+    ) -> Self;
 
     /// Bootstrap the sink
     ///
@@ -77,7 +81,11 @@ mod tests {
     impl Sink for TestSink {
         type Config = Option<()>;
 
-        fn new(config: Option<()>, _stats: InputQueueScope) -> Self {
+        fn new(
+            config: Option<()>,
+            _schemas: &[crate::settings::Schema],
+            _stats: InputQueueScope,
+        ) -> Self {
             Self { config }
         }