Non-data fields - Metadata fields to Kafka Connector (#765)

ArroyoSystems · Oct 30, 2024 · 28360fe · 28360fe
1 parent 790b47d
commit 28360fe
Show file tree

Hide file tree

Showing 40 changed files with 617 additions and 47 deletions.
diff --git a/crates/arroyo-api/src/pipelines.rs b/crates/arroyo-api/src/pipelines.rs
@@ -159,6 +159,7 @@ async fn compile_sql<'a>(
                     .unwrap_or(json!({})),
                 &table.config,
                 Some(&table.schema),
+                None,
             )
             .map_err(log_and_map)?;
 

diff --git a/crates/arroyo-connectors/src/blackhole/mod.rs b/crates/arroyo-connectors/src/blackhole/mod.rs
@@ -1,5 +1,6 @@
 use crate::blackhole::operator::BlackholeSinkFunc;
 use anyhow::anyhow;
+use arrow::datatypes::DataType;
 use arroyo_operator::connector::{Connection, Connector};
 use arroyo_operator::operator::OperatorNode;
 use arroyo_rpc::api_types::connections::{
@@ -78,8 +79,9 @@ impl Connector for BlackholeConnector {
         _options: &mut HashMap<String, String>,
         schema: Option<&ConnectionSchema>,
         _profile: Option<&ConnectionProfile>,
+        _metadata_fields: Option<HashMap<String, (String, DataType)>>,
     ) -> anyhow::Result<Connection> {
-        self.from_config(None, name, EmptyConfig {}, EmptyConfig {}, schema)
+        self.from_config(None, name, EmptyConfig {}, EmptyConfig {}, schema, None)
     }
 
     fn from_config(
@@ -89,6 +91,7 @@ impl Connector for BlackholeConnector {
         config: Self::ProfileT,
         table: Self::TableT,
         s: Option<&ConnectionSchema>,
+        _metadata_fields: Option<HashMap<String, (String, DataType)>>,
     ) -> anyhow::Result<Connection> {
         let description = "Blackhole".to_string();
 
@@ -99,6 +102,7 @@ impl Connector for BlackholeConnector {
             format: None,
             bad_data: None,
             framing: None,
+            additional_fields: None,
         };
 
         Ok(Connection {

diff --git a/crates/arroyo-connectors/src/confluent/mod.rs b/crates/arroyo-connectors/src/confluent/mod.rs
@@ -3,6 +3,7 @@ use crate::kafka::{
 };
 use crate::{kafka, pull_opt};
 use anyhow::anyhow;
+use arrow::datatypes::DataType;
 use arroyo_operator::connector::{Connection, Connector};
 use arroyo_operator::operator::OperatorNode;
 use arroyo_rpc::api_types::connections::{
@@ -161,6 +162,7 @@ impl Connector for ConfluentConnector {
         options: &mut HashMap<String, String>,
         schema: Option<&ConnectionSchema>,
         profile: Option<&ConnectionProfile>,
+        _metadata_fields: Option<HashMap<String, (String, DataType)>>,
     ) -> anyhow::Result<Connection> {
         let connection = profile
             .map(|p| {
@@ -172,7 +174,7 @@ impl Connector for ConfluentConnector {
 
         let table = KafkaConnector::table_from_options(options)?;
 
-        self.from_config(None, name, connection, table, schema)
+        self.from_config(None, name, connection, table, schema, None)
     }
 
     fn from_config(
@@ -182,11 +184,12 @@ impl Connector for ConfluentConnector {
         config: Self::ProfileT,
         mut table: Self::TableT,
         schema: Option<&ConnectionSchema>,
+        _metadata_fields: Option<HashMap<String, (String, DataType)>>,
     ) -> anyhow::Result<Connection> {
         table
             .client_configs
             .insert("client.id".to_string(), CLIENT_ID.to_string());
-        KafkaConnector {}.from_config(id, name, config.into(), table, schema)
+        KafkaConnector {}.from_config(id, name, config.into(), table, schema, None)
     }
 
     fn make_operator(

diff --git a/crates/arroyo-connectors/src/filesystem/delta.rs b/crates/arroyo-connectors/src/filesystem/delta.rs
@@ -1,4 +1,5 @@
 use anyhow::{anyhow, bail};
+use arrow::datatypes::DataType;
 use arroyo_operator::connector::Connection;
 use arroyo_storage::BackendConfig;
 use std::collections::HashMap;
@@ -77,6 +78,7 @@ impl Connector for DeltaLakeConnector {
         config: Self::ProfileT,
         table: Self::TableT,
         schema: Option<&ConnectionSchema>,
+        _metadata_fields: Option<HashMap<String, (String, DataType)>>,
     ) -> anyhow::Result<arroyo_operator::connector::Connection> {
         let TableType::Sink {
             write_path,
@@ -123,6 +125,7 @@ impl Connector for DeltaLakeConnector {
             format: Some(format),
             bad_data: schema.bad_data.clone(),
             framing: schema.framing.clone(),
+            additional_fields: None,
         };
 
         Ok(Connection {
@@ -142,10 +145,11 @@ impl Connector for DeltaLakeConnector {
         options: &mut HashMap<String, String>,
         schema: Option<&ConnectionSchema>,
         _profile: Option<&ConnectionProfile>,
+        _metadata_fields: Option<HashMap<String, (String, DataType)>>,
     ) -> anyhow::Result<Connection> {
         let table = file_system_sink_from_options(options, schema, CommitStyle::DeltaLake)?;
 
-        self.from_config(None, name, EmptyConfig {}, table, schema)
+        self.from_config(None, name, EmptyConfig {}, table, schema, None)
     }
 
     fn make_operator(

diff --git a/crates/arroyo-connectors/src/filesystem/mod.rs b/crates/arroyo-connectors/src/filesystem/mod.rs
@@ -3,6 +3,7 @@ mod sink;
 mod source;
 
 use anyhow::{anyhow, bail, Result};
+use arrow::datatypes::DataType;
 use arroyo_storage::BackendConfig;
 use regex::Regex;
 use std::collections::HashMap;
@@ -114,6 +115,7 @@ impl Connector for FileSystemConnector {
         config: Self::ProfileT,
         table: Self::TableT,
         schema: Option<&ConnectionSchema>,
+        _metadata_fields: Option<HashMap<String, (String, DataType)>>,
     ) -> anyhow::Result<Connection> {
         let (description, connection_type) = match table.table_type {
             TableType::Source { .. } => ("FileSystem".to_string(), ConnectionType::Source),
@@ -168,6 +170,7 @@ impl Connector for FileSystemConnector {
             format: Some(format),
             bad_data: schema.bad_data.clone(),
             framing: schema.framing.clone(),
+            additional_fields: None,
         };
 
         Ok(Connection {
@@ -187,6 +190,7 @@ impl Connector for FileSystemConnector {
         options: &mut HashMap<String, String>,
         schema: Option<&ConnectionSchema>,
         _profile: Option<&ConnectionProfile>,
+        _metadata_fields: Option<HashMap<String, (String, DataType)>>,
     ) -> anyhow::Result<Connection> {
         match options.remove("type") {
             Some(t) if t == "source" => {
@@ -210,12 +214,13 @@ impl Connector for FileSystemConnector {
                         },
                     },
                     schema,
+                    None,
                 )
             }
             Some(t) if t == "sink" => {
                 let table = file_system_sink_from_options(options, schema, CommitStyle::Direct)?;
 
-                self.from_config(None, name, EmptyConfig {}, table, schema)
+                self.from_config(None, name, EmptyConfig {}, table, schema, None)
             }
             Some(t) => bail!("unknown type: {}", t),
             None => bail!("must have type set"),

diff --git a/crates/arroyo-connectors/src/filesystem/source.rs b/crates/arroyo-connectors/src/filesystem/source.rs
@@ -367,7 +367,7 @@ impl FileSystemSourceFunc {
                 line = line_reader.next() => {
                     match line.transpose()? {
                         Some(line) => {
-                            ctx.deserialize_slice(line.as_bytes(), SystemTime::now()).await?;
+                            ctx.deserialize_slice(line.as_bytes(), SystemTime::now(), None).await?;
                             records_read += 1;
                             if ctx.should_flush() {
                                 ctx.flush_buffer().await?;

diff --git a/crates/arroyo-connectors/src/fluvio/mod.rs b/crates/arroyo-connectors/src/fluvio/mod.rs
@@ -1,4 +1,5 @@
 use anyhow::{anyhow, bail};
+use arrow::datatypes::DataType;
 use arroyo_formats::ser::ArrowSerializer;
 use arroyo_operator::connector::{Connection, Connector};
 use arroyo_operator::operator::OperatorNode;
@@ -88,6 +89,7 @@ impl Connector for FluvioConnector {
         options: &mut HashMap<String, String>,
         schema: Option<&ConnectionSchema>,
         _profile: Option<&ConnectionProfile>,
+        _metadata_fields: Option<HashMap<String, (String, DataType)>>,
     ) -> anyhow::Result<Connection> {
         let endpoint = options.remove("endpoint");
         let topic = pull_opt("topic", options)?;
@@ -116,7 +118,7 @@ impl Connector for FluvioConnector {
             type_: table_type,
         };
 
-        Self::from_config(self, None, name, EmptyConfig {}, table, schema)
+        Self::from_config(self, None, name, EmptyConfig {}, table, schema, None)
     }
 
     fn from_config(
@@ -126,6 +128,7 @@ impl Connector for FluvioConnector {
         config: EmptyConfig,
         table: FluvioTable,
         schema: Option<&ConnectionSchema>,
+        _metadata_fields: Option<HashMap<String, (String, DataType)>>,
     ) -> anyhow::Result<Connection> {
         let (typ, desc) = match table.type_ {
             TableType::Source { .. } => (
@@ -154,6 +157,7 @@ impl Connector for FluvioConnector {
             format: Some(format),
             bad_data: schema.bad_data.clone(),
             framing: schema.framing.clone(),
+            additional_fields: None,
         };
 
         Ok(Connection {

diff --git a/crates/arroyo-connectors/src/fluvio/source.rs b/crates/arroyo-connectors/src/fluvio/source.rs
@@ -166,7 +166,7 @@ impl FluvioSourceFunc {
                     match message {
                         Some((_, Ok(msg))) => {
                             let timestamp = from_millis(msg.timestamp().max(0) as u64);
-                            ctx.deserialize_slice(msg.value(), timestamp).await?;
+                            ctx.deserialize_slice(msg.value(), timestamp, None).await?;
 
                             if ctx.should_flush() {
                                 ctx.flush_buffer().await?;

diff --git a/crates/arroyo-connectors/src/impulse/mod.rs b/crates/arroyo-connectors/src/impulse/mod.rs
@@ -1,6 +1,7 @@
 mod operator;
 
 use anyhow::{anyhow, bail};
+use arrow::datatypes::DataType;
 use arroyo_operator::connector::{Connection, Connector};
 use arroyo_operator::operator::OperatorNode;
 use arroyo_rpc::api_types::connections::FieldType::Primitive;
@@ -101,6 +102,7 @@ impl Connector for ImpulseConnector {
         options: &mut HashMap<String, String>,
         schema: Option<&ConnectionSchema>,
         _profile: Option<&ConnectionProfile>,
+        _metadata_fields: Option<HashMap<String, (String, DataType)>>,
     ) -> anyhow::Result<Connection> {
         let event_rate = f64::from_str(&pull_opt("event_rate", options)?)
             .map_err(|_| anyhow!("invalid value for event_rate; expected float"))?;
@@ -134,6 +136,7 @@ impl Connector for ImpulseConnector {
                 message_count,
             },
             None,
+            None,
         )
     }
 
@@ -144,6 +147,7 @@ impl Connector for ImpulseConnector {
         config: Self::ProfileT,
         table: Self::TableT,
         _: Option<&ConnectionSchema>,
+        _metadata_fields: Option<HashMap<String, (String, DataType)>>,
     ) -> anyhow::Result<Connection> {
         let description = format!(
             "{}Impulse<{} eps{}>",
@@ -166,6 +170,7 @@ impl Connector for ImpulseConnector {
             format: None,
             bad_data: None,
             framing: None,
+            additional_fields: None,
         };
 
         Ok(Connection {

diff --git a/crates/arroyo-connectors/src/kafka/mod.rs b/crates/arroyo-connectors/src/kafka/mod.rs
@@ -1,4 +1,5 @@
 use anyhow::{anyhow, bail};
+use arrow::datatypes::DataType;
 use arroyo_formats::de::ArrowDeserializer;
 use arroyo_formats::ser::ArrowSerializer;
 use arroyo_operator::connector::Connection;
@@ -188,6 +189,7 @@ impl Connector for KafkaConnector {
         config: KafkaConfig,
         table: KafkaTable,
         schema: Option<&ConnectionSchema>,
+        metadata_fields: Option<HashMap<String, (String, DataType)>>,
     ) -> anyhow::Result<Connection> {
         let (typ, desc) = match table.type_ {
             TableType::Source { .. } => (
@@ -207,13 +209,21 @@ impl Connector for KafkaConnector {
             .map(|t| t.to_owned())
             .ok_or_else(|| anyhow!("'format' must be set for Kafka connection"))?;
 
+        let metadata_fields = metadata_fields.map(|fields| {
+            fields
+                .into_iter()
+                .map(|(k, (v, _))| (k, v))
+                .collect::<HashMap<String, String>>()
+        });
+
         let config = OperatorConfig {
             connection: serde_json::to_value(config).unwrap(),
             table: serde_json::to_value(table).unwrap(),
             rate_limit: None,
             format: Some(format),
             bad_data: schema.bad_data.clone(),
             framing: schema.framing.clone(),
+            additional_fields: metadata_fields,
         };
 
         Ok(Connection {
@@ -312,6 +322,7 @@ impl Connector for KafkaConnector {
         options: &mut HashMap<String, String>,
         schema: Option<&ConnectionSchema>,
         profile: Option<&ConnectionProfile>,
+        metadata_fields: Option<HashMap<String, (String, DataType)>>,
     ) -> anyhow::Result<Connection> {
         let connection = profile
             .map(|p| {
@@ -323,7 +334,37 @@ impl Connector for KafkaConnector {
 
         let table = Self::table_from_options(options)?;
 
-        Self::from_config(self, None, name, connection, table, schema)
+        let allowed_metadata_udf_args: HashMap<&str, DataType> = [
+            ("offset_id", DataType::Int64),
+            ("partition", DataType::Int32),
+            ("topic", DataType::Utf8),
+        ]
+        .iter()
+        .cloned()
+        .collect();
+
+        if let Some(fields) = &metadata_fields {
+            for (field_name, data_type) in fields.values() {
+                match allowed_metadata_udf_args.get(field_name.as_str()) {
+                    Some(expected_type) => {
+                        if expected_type != data_type {
+                            return Err(anyhow!(
+                                "Invalid datatype for metadata field '{}': expected '{:?}', found '{:?}'",
+                                field_name, expected_type, data_type
+                            ));
+                        }
+                    }
+                    None => {
+                        return Err(anyhow!(
+                            "Invalid metadata field name for Kafka connector: '{}'",
+                            field_name
+                        ));
+                    }
+                }
+            }
+        }
+
+        Self::from_config(self, None, name, connection, table, schema, metadata_fields)
     }
 
     fn make_operator(
@@ -383,6 +424,7 @@ impl Connector for KafkaConnector {
                             .unwrap_or(u32::MAX),
                     )
                     .unwrap(),
+                    metadata_fields: config.additional_fields,
                 })))
             }
             TableType::Sink {
@@ -622,7 +664,7 @@ impl KafkaTester {
                     let mut builders = aschema.builders();
 
                     let mut error = deserializer
-                        .deserialize_slice(&mut builders, &msg, SystemTime::now())
+                        .deserialize_slice(&mut builders, &msg, SystemTime::now(), None)
                         .await
                         .into_iter()
                         .next();
@@ -644,7 +686,7 @@ impl KafkaTester {
                     let mut builders = aschema.builders();
 
                     let mut error = deserializer
-                        .deserialize_slice(&mut builders, &msg, SystemTime::now())
+                        .deserialize_slice(&mut builders, &msg, SystemTime::now(), None)
                         .await
                         .into_iter()
                         .next();
@@ -678,7 +720,7 @@ impl KafkaTester {
                 let mut builders = aschema.builders();
 
                 let mut error = deserializer
-                    .deserialize_slice(&mut builders, &msg, SystemTime::now())
+                    .deserialize_slice(&mut builders, &msg, SystemTime::now(), None)
                     .await
                     .into_iter()
                     .next();