Add control for deserialization error behavior

Adds a config option 'bad_data' to control whether deserialization errors result in dropping the data or failing the job. The logic is mostly contained in `source_collector.rs`, and it uses a new `RateLimiter` struct to limit the logging frequency and user error reporting.
ArroyoSystems · Dec 13, 2023 · 60b62d1 · 60b62d1
1 parent f6a9964
commit 60b62d1
Show file tree

Hide file tree

Showing 36 changed files with 288 additions and 99 deletions.
diff --git a/arroyo-api/src/lib.rs b/arroyo-api/src/lib.rs
@@ -219,6 +219,7 @@ pub(crate) fn to_micros(dt: OffsetDateTime) -> u64 {
         UdfPost,
         GlobalUdf,
         GlobalUdfCollection,
+        BadData,
     )),
     tags(
         (name = "ping", description = "Ping endpoint"),

diff --git a/arroyo-connectors/src/blackhole.rs b/arroyo-connectors/src/blackhole.rs
@@ -96,6 +96,7 @@ impl Connector for BlackholeConnector {
             table: serde_json::to_value(table).unwrap(),
             rate_limit: None,
             format: None,
+            bad_data: None,
             framing: None,
         };
 

diff --git a/arroyo-connectors/src/delta.rs b/arroyo-connectors/src/delta.rs
@@ -130,6 +130,7 @@ impl Connector for DeltaLakeConnector {
             table: serde_json::to_value(table).unwrap(),
             rate_limit: None,
             format: Some(format),
+            bad_data: schema.invalid_data_behavior.clone(),
             framing: schema.framing.clone(),
         };
 

diff --git a/arroyo-connectors/src/filesystem.rs b/arroyo-connectors/src/filesystem.rs
@@ -148,6 +148,7 @@ impl Connector for FileSystemConnector {
             table: serde_json::to_value(table).unwrap(),
             rate_limit: None,
             format: Some(format),
+            bad_data: schema.invalid_data_behavior.clone(),
             framing: schema.framing.clone(),
         };
 

diff --git a/arroyo-connectors/src/fluvio.rs b/arroyo-connectors/src/fluvio.rs
@@ -150,6 +150,7 @@ impl Connector for FluvioConnector {
             table: serde_json::to_value(table).unwrap(),
             rate_limit: None,
             format: Some(format),
+            bad_data: schema.invalid_data_behavior.clone(),
             framing: schema.framing.clone(),
         };
 

diff --git a/arroyo-connectors/src/impulse.rs b/arroyo-connectors/src/impulse.rs
@@ -22,6 +22,7 @@ pub fn impulse_schema() -> ConnectionSchema {
     ConnectionSchema {
         format: None,
         framing: None,
+        invalid_data_behavior: None,
         struct_name: Some("arroyo_types::ImpulseEvent".to_string()),
         fields: vec![
             source_field("counter", Primitive(PrimitiveType::UInt64)),
@@ -161,6 +162,7 @@ impl Connector for ImpulseConnector {
             table: serde_json::to_value(table).unwrap(),
             rate_limit: None,
             format: None,
+            bad_data: None,
             framing: None,
         };
 

diff --git a/arroyo-connectors/src/kafka.rs b/arroyo-connectors/src/kafka.rs
@@ -99,6 +99,7 @@ impl Connector for KafkaConnector {
             table: serde_json::to_value(table).unwrap(),
             rate_limit: None,
             format: Some(format),
+            bad_data: schema.invalid_data_behavior.clone(),
             framing: schema.framing.clone(),
         };
 

diff --git a/arroyo-connectors/src/kinesis.rs b/arroyo-connectors/src/kinesis.rs
@@ -109,6 +109,7 @@ impl Connector for KinesisConnector {
             table: serde_json::to_value(table).unwrap(),
             rate_limit: None,
             format: Some(format),
+            bad_data: schema.invalid_data_behavior.clone(),
             framing: schema.framing.clone(),
         };
 

diff --git a/arroyo-connectors/src/nexmark.rs b/arroyo-connectors/src/nexmark.rs
@@ -23,6 +23,7 @@ pub fn nexmark_schema() -> ConnectionSchema {
     use arroyo_rpc::api_types::connections::PrimitiveType::*;
     ConnectionSchema {
         format: None,
+        invalid_data_behavior: None,
         framing: None,
         struct_name: Some("arroyo_types::nexmark::Event".to_string()),
         fields: vec![
@@ -207,6 +208,7 @@ impl Connector for NexmarkConnector {
             table: serde_json::to_value(table).unwrap(),
             rate_limit: None,
             format: None,
+            bad_data: None,
             framing: None,
         };
 

diff --git a/arroyo-connectors/src/polling_http.rs b/arroyo-connectors/src/polling_http.rs
@@ -218,6 +218,7 @@ impl Connector for PollingHTTPConnector {
             table: serde_json::to_value(table).unwrap(),
             rate_limit: None,
             format: Some(format),
+            bad_data: schema.invalid_data_behavior.clone(),
             framing: schema.framing.clone(),
         };
 

diff --git a/arroyo-connectors/src/redis.rs b/arroyo-connectors/src/redis.rs
@@ -373,6 +373,7 @@ impl Connector for RedisConnector {
             table: serde_json::to_value(table).unwrap(),
             rate_limit: None,
             format: Some(format),
+            bad_data: schema.invalid_data_behavior.clone(),
             framing: schema.framing.clone(),
         };
 

diff --git a/arroyo-connectors/src/single_file.rs b/arroyo-connectors/src/single_file.rs
@@ -105,6 +105,7 @@ impl Connector for SingleFileConnector {
             table: serde_json::to_value(table).unwrap(),
             rate_limit: None,
             format: Some(format),
+            bad_data: schema.invalid_data_behavior.clone(),
             framing: schema.framing.clone(),
         };
 

diff --git a/arroyo-connectors/src/sse.rs b/arroyo-connectors/src/sse.rs
@@ -102,6 +102,7 @@ impl Connector for SSEConnector {
             table: serde_json::to_value(table).unwrap(),
             rate_limit: None,
             format: Some(format),
+            bad_data: schema.invalid_data_behavior.clone(),
             framing: schema.framing.clone(),
         };
 

diff --git a/arroyo-connectors/src/webhook.rs b/arroyo-connectors/src/webhook.rs
@@ -158,6 +158,7 @@ impl Connector for WebhookConnector {
             table: serde_json::to_value(table).unwrap(),
             rate_limit: None,
             format: Some(format),
+            bad_data: schema.invalid_data_behavior.clone(),
             framing: schema.framing.clone(),
         };
 

diff --git a/arroyo-connectors/src/websocket.rs b/arroyo-connectors/src/websocket.rs
@@ -246,6 +246,7 @@ impl Connector for WebsocketConnector {
             table: serde_json::to_value(table).unwrap(),
             rate_limit: None,
             format: Some(format),
+            bad_data: schema.invalid_data_behavior.clone(),
             framing: schema.framing.clone(),
         };
 

diff --git a/arroyo-formats/src/avro.rs b/arroyo-formats/src/avro.rs
@@ -4,7 +4,7 @@ use apache_avro::{from_avro_datum, Reader, Schema, Writer};
 use arrow::datatypes::{DataType, Field, Fields, TimeUnit};
 use arroyo_rpc::formats::AvroFormat;
 use arroyo_rpc::schema_resolver::SchemaResolver;
-use arroyo_types::UserError;
+use arroyo_types::SourceError;
 use serde::de::DeserializeOwned;
 use serde_json::{json, Value as JsonValue};
 use std::collections::HashMap;
@@ -17,7 +17,7 @@ pub async fn deserialize_slice_avro<'a, T: DeserializeOwned>(
     schema_registry: Arc<Mutex<HashMap<u32, Schema>>>,
     resolver: Arc<dyn SchemaResolver + Sync>,
     mut msg: &'a [u8],
-) -> Result<impl Iterator<Item = Result<T, UserError>> + 'a, String> {
+) -> Result<impl Iterator<Item = Result<T, SourceError>> + 'a, String> {
     let id = if format.confluent_schema_registry {
         let magic_byte = msg[0];
         if magic_byte != 0 {
@@ -72,10 +72,7 @@ pub async fn deserialize_slice_avro<'a, T: DeserializeOwned>(
     let into_json = format.into_unstructured_json;
     Ok(messages.into_iter().map(move |record| {
         let value = record.map_err(|e| {
-            UserError::new(
-                "Deserialization failed",
-                format!("Failed to deserialize from avro: {:?}", e),
-            )
+            SourceError::bad_data(format!("Failed to deserialize from avro: {:?}", e))
         })?;
 
         if into_json {
@@ -84,10 +81,10 @@ pub async fn deserialize_slice_avro<'a, T: DeserializeOwned>(
             // for now round-trip through json in order to handle unsupported avro features
             // as that allows us to rely on raw json deserialization
             serde_json::from_value(avro_to_json(value)).map_err(|e| {
-                UserError::new(
-                    "Deserialization failed",
-                    format!("Failed to convert avro message into struct type: {:?}", e),
-                )
+                SourceError::bad_data(format!(
+                    "Failed to convert avro message into struct type: {:?}",
+                    e
+                ))
             })
         }
     }))

diff --git a/arroyo-formats/src/lib.rs b/arroyo-formats/src/lib.rs
@@ -6,7 +6,7 @@ use arrow_array::cast::AsArray;
 use arrow_array::{RecordBatch, StringArray};
 use arroyo_rpc::formats::{AvroFormat, Format, Framing, FramingMethod};
 use arroyo_rpc::schema_resolver::{FailingSchemaResolver, FixedSchemaResolver, SchemaResolver};
-use arroyo_types::{Data, Debezium, RawJson, UserError};
+use arroyo_types::{Data, Debezium, RawJson, SourceError};
 use serde::de::DeserializeOwned;
 use serde::{Deserialize, Deserializer, Serialize};
 use serde_json::{json, Value};
@@ -290,7 +290,7 @@ impl<T: SchemaData> DataDeserializer<T> {
     pub async fn deserialize_slice<'a>(
         &mut self,
         msg: &'a [u8],
-    ) -> impl Iterator<Item = Result<T, UserError>> + 'a + Send {
+    ) -> impl Iterator<Item = Result<T, SourceError>> + 'a + Send {
         match &*self.format {
             Format::Avro(avro) => {
                 let schema_registry = self.schema_registry.clone();
@@ -300,17 +300,21 @@ impl<T: SchemaData> DataDeserializer<T> {
                 {
                     Ok(iter) => Box::new(iter),
                     Err(e) => Box::new(
-                        vec![Err(UserError::new("Avro deserialization failed", e))].into_iter(),
+                        vec![Err(SourceError::bad_data(format!(
+                            "Avro deserialization failed: {}",
+                            e
+                        )))]
+                        .into_iter(),
                     )
-                        as Box<dyn Iterator<Item = Result<T, UserError>> + Send>,
+                        as Box<dyn Iterator<Item = Result<T, SourceError>> + Send>,
                 }
             }
             _ => {
                 let new_self = self.clone();
                 Box::new(
                     FramingIterator::new(self.framing.clone(), msg)
                         .map(move |t| new_self.deserialize_single(t)),
-                ) as Box<dyn Iterator<Item = Result<T, UserError>> + Send>
+                ) as Box<dyn Iterator<Item = Result<T, SourceError>> + Send>
             }
         }
     }
@@ -319,19 +323,14 @@ impl<T: SchemaData> DataDeserializer<T> {
         self.format.clone()
     }
 
-    pub fn deserialize_single(&self, msg: &[u8]) -> Result<T, UserError> {
+    pub fn deserialize_single(&self, msg: &[u8]) -> Result<T, SourceError> {
         match &*self.format {
             Format::Json(json) => json::deserialize_slice_json(json, msg),
             Format::Avro(_) => unreachable!("avro should be handled by here"),
             Format::Parquet(_) => todo!("parquet is not supported as an input format"),
             Format::RawString(_) => deserialize_raw_string(msg),
         }
-        .map_err(|e| {
-            UserError::new(
-                "Deserialization failed",
-                format!("Failed to deserialize: {}", e),
-            )
-        })
+        .map_err(|e| SourceError::bad_data(format!("Failed to deserialize: {:?}", e)))
     }
 }
 

diff --git a/arroyo-rpc/src/api_types/connections.rs b/arroyo-rpc/src/api_types/connections.rs
@@ -1,4 +1,4 @@
-use crate::formats::{Format, Framing};
+use crate::formats::{BadData, Format, Framing};
 use anyhow::bail;
 use serde::{Deserialize, Serialize};
 use std::fmt::{Display, Formatter};
@@ -128,6 +128,7 @@ pub enum SchemaDefinition {
 #[serde(rename_all = "camelCase")]
 pub struct ConnectionSchema {
     pub format: Option<Format>,
+    pub invalid_data_behavior: Option<BadData>,
     pub framing: Option<Framing>,
     pub struct_name: Option<String>,
     pub fields: Vec<SourceField>,
@@ -138,6 +139,7 @@ pub struct ConnectionSchema {
 impl ConnectionSchema {
     pub fn try_new(
         format: Option<Format>,
+        invalid_data_behavior: Option<BadData>,
         framing: Option<Framing>,
         struct_name: Option<String>,
         fields: Vec<SourceField>,
@@ -146,6 +148,7 @@ impl ConnectionSchema {
     ) -> anyhow::Result<Self> {
         let s = ConnectionSchema {
             format,
+            invalid_data_behavior,
             framing,
             struct_name,
             fields,

diff --git a/arroyo-rpc/src/formats.rs b/arroyo-rpc/src/formats.rs
@@ -260,6 +260,29 @@ impl Format {
     }
 }
 
+#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, Hash, PartialOrd, ToSchema)]
+#[serde(rename_all = "snake_case")]
+pub enum BadData {
+    Drop,
+    Fail,
+}
+
+impl BadData {
+    pub fn from_opts(opts: &mut HashMap<String, String>) -> Result<Option<Self>, String> {
+        let Some(method) = opts.remove("invalid_data_behavior") else {
+            return Ok(None);
+        };
+
+        let method = match method.as_str() {
+            "drop" => BadData::Drop,
+            "fail" => BadData::Fail,
+            f => return Err(format!("Unknown invalid data behavior '{}'", f)),
+        };
+
+        Ok(Some(method))
+    }
+}
+
 #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, Hash, PartialOrd, ToSchema)]
 #[serde(rename_all = "camelCase")]
 pub struct Framing {

diff --git a/arroyo-rpc/src/lib.rs b/arroyo-rpc/src/lib.rs
@@ -8,7 +8,7 @@ use std::collections::HashMap;
 use std::{fs, time::SystemTime};
 
 use crate::api_types::connections::PrimitiveType;
-use crate::formats::{Format, Framing};
+use crate::formats::{BadData, Format, Framing};
 use crate::grpc::{LoadCompactedDataReq, SubtaskCheckpointMetadata};
 use arroyo_types::CheckpointBarrier;
 use grpc::{StopMode, TaskCheckpointEventType};
@@ -165,6 +165,7 @@ pub struct OperatorConfig {
     pub connection: Value,
     pub table: Value,
     pub format: Option<Format>,
+    pub bad_data: Option<BadData>,
     pub framing: Option<Framing>,
     pub rate_limit: Option<RateLimit>,
 }
@@ -175,6 +176,7 @@ impl Default for OperatorConfig {
             connection: serde_json::from_str("{}").unwrap(),
             table: serde_json::from_str("{}").unwrap(),
             format: None,
+            bad_data: None,
             framing: None,
             rate_limit: None,
         }

diff --git a/arroyo-sql-macro/src/connectors.rs b/arroyo-sql-macro/src/connectors.rs
@@ -80,6 +80,7 @@ pub fn get_json_schema_source() -> Result<Connection> {
         })),
         None,
         None,
+        None,
         struct_fields
             .into_iter()
             .map(|field| field.try_into().unwrap())
@@ -185,6 +186,7 @@ pub fn get_avro_source() -> Result<Connection> {
         Some(Format::Avro(format)),
         None,
         None,
+        None,
         struct_fields
             .into_iter()
             .map(|field| field.try_into().unwrap())

diff --git a/arroyo-sql/src/lib.rs b/arroyo-sql/src/lib.rs
@@ -684,6 +684,7 @@ pub fn get_test_expression(
     let struct_def = test_struct_def();
     let schema = ConnectionSchema {
         format: Some(Format::Json(JsonFormat::default())),
+        invalid_data_behavior: None,
         framing: None,
         struct_name: struct_def.name.clone(),
         fields: struct_def

diff --git a/arroyo-sql/src/tables.rs b/arroyo-sql/src/tables.rs
@@ -8,7 +8,7 @@ use arroyo_datastream::{ConnectorOp, Operator};
 use arroyo_rpc::api_types::connections::{
     ConnectionProfile, ConnectionSchema, ConnectionType, SchemaDefinition, SourceField,
 };
-use arroyo_rpc::formats::{Format, Framing};
+use arroyo_rpc::formats::{BadData, Format, Framing};
 use datafusion::sql::sqlparser::ast::Query;
 use datafusion::{
     optimizer::{analyzer::Analyzer, optimizer::Optimizer, OptimizerContext},
@@ -190,6 +190,9 @@ impl ConnectorTable {
 
         let format = Format::from_opts(options).map_err(|e| anyhow!("invalid format: '{e}'"))?;
 
+        let invalid_data_behavior = BadData::from_opts(options)
+            .map_err(|e| anyhow!("Invalid invalid_data_behavior: '{e}'"))?;
+
         let framing = Framing::from_opts(options).map_err(|e| anyhow!("invalid framing: '{e}'"))?;
 
         let schema_fields: Result<Vec<SourceField>> = fields
@@ -209,6 +212,7 @@ impl ConnectorTable {
 
         let schema = ConnectionSchema::try_new(
             format,
+            invalid_data_behavior,
             framing,
             None,
             schema_fields?,