fix: hierarchical json flattening restriction

get the level of hierarchy from the json
perform generic flattening only if level of nesting is <=4

Author: nikhilsinhaparseable

src/cli.rs

@@ -529,7 +529,9 @@ impl FromArgMatches for Cli {
         self.kafka_host = m.get_one::<String>(Self::KAFKA_HOST).cloned();
         self.kafka_group = m.get_one::<String>(Self::KAFKA_GROUP).cloned();
         self.kafka_client_id = m.get_one::<String>(Self::KAFKA_CLIENT_ID).cloned();
-        self.kafka_security_protocol = m.get_one::<SslProtocol>(Self::KAFKA_SECURITY_PROTOCOL).cloned();
+        self.kafka_security_protocol = m
+            .get_one::<SslProtocol>(Self::KAFKA_SECURITY_PROTOCOL)
+            .cloned();
         self.kafka_partitions = m.get_one::<String>(Self::KAFKA_PARTITIONS).cloned();
 
         self.tls_cert_path = m.get_one::<PathBuf>(Self::TLS_CERT).cloned();

src/kafka.rs

@@ -125,7 +125,7 @@ pub enum KafkaError {
 fn setup_consumer() -> Result<(StreamConsumer, Vec<String>), KafkaError> {
     if let Some(topics) = &CONFIG.parseable.kafka_topics {
         // topics can be a comma separated list of topics to subscribe to
-        let topics = topics.split(",").map(|v| v.to_owned()).collect_vec();
+        let topics = topics.split(',').map(|v| v.to_owned()).collect_vec();
 
         let host = if CONFIG.parseable.kafka_host.is_some() {
             CONFIG.parseable.kafka_host.as_ref()
@@ -162,8 +162,8 @@ fn setup_consumer() -> Result<(StreamConsumer, Vec<String>), KafkaError> {
             // partitions is a comma separated pairs of topic:partitions
             let mut topic_partition_pairs = Vec::new();
             let mut set = true;
-            for vals in vals_raw.split(",") {
-                let intermediate = vals.split(":").collect_vec();
+            for vals in vals_raw.split(',') {
+                let intermediate = vals.split(':').collect_vec();
                 if intermediate.len() != 2 {
                     warn!(
                         "Value for P_KAFKA_PARTITIONS is incorrect! Skipping setting partitions!"

src/utils/json/flatten.rs

@@ -19,6 +19,7 @@
 use std::collections::BTreeMap;
 use std::num::NonZeroU32;
 
+use anyhow::anyhow;
 use chrono::{DateTime, Duration, Utc};
 use serde_json::map::Map;
 use serde_json::value::Value;
@@ -273,52 +274,83 @@ pub fn flatten_array_objects(
 /// Recursively flattens a JSON value.
 /// - If the value is an array, it flattens all elements of the array.
 /// - If the value is an object, it flattens all nested objects and arrays.
+/// - If the JSON value is heavily nested (with more than 4 levels of hierarchy), returns error
 /// - Otherwise, it returns the value itself in a vector.
 ///
 /// Examples:
 /// 1. `{"a": 1}` ~> `[{"a": 1}]`
 /// 2. `[{"a": 1}, {"b": 2}]` ~> `[{"a": 1}, {"b": 2}]`
 /// 3. `[{"a": [{"b": 1}, {"c": 2}]}]` ~> `[{"a": {"b": 1)}}, {"a": {"c": 2)}}]`
-/// 3. `{"a": [{"b": 1}, {"c": 2}], "d": {"e": 4}}` ~> `[{"a": {"b":1}, "d": {"e":4}}, {"a": {"c":2}, "d": {"e":4}}]`
-pub fn flatten_json(value: &Value) -> Vec<Value> {
+/// 4. `{"a": [{"b": 1}, {"c": 2}], "d": {"e": 4}}` ~> `[{"a": {"b":1}, "d": {"e":4}}, {"a": {"c":2}, "d": {"e":4}}]`
+/// 5. `{"a":{"b":{"c":{"d":{"e":["a","b"]}}}}}` ~> returns error - heavily nested, cannot flatten this JSON
+pub fn flatten_json(value: &Value) -> Result<Vec<Value>, anyhow::Error> {
+    if has_more_than_four_levels(value, 1) {
+        return Err(anyhow!("heavily nested, cannot flatten this JSON"));
+    }
+
     match value {
-        Value::Array(arr) => arr.iter().flat_map(flatten_json).collect(),
-        Value::Object(map) => map
+        Value::Array(arr) => Ok(arr
             .iter()
-            .fold(vec![Map::new()], |results, (key, val)| match val {
-                Value::Array(arr) => arr
-                    .iter()
-                    .flat_map(flatten_json)
-                    .flat_map(|flattened_item| {
-                        results.iter().map(move |result| {
-                            let mut new_obj = result.clone();
-                            new_obj.insert(key.clone(), flattened_item.clone());
-                            new_obj
+            .flat_map(|flatten_item| flatten_json(flatten_item).unwrap_or_default())
+            .collect()),
+        Value::Object(map) => {
+            let results = map
+                .iter()
+                .fold(vec![Map::new()], |results, (key, val)| match val {
+                    Value::Array(arr) => arr
+                        .iter()
+                        .flat_map(|flatten_item| flatten_json(flatten_item).unwrap_or_default())
+                        .flat_map(|flattened_item| {
+                            results.iter().map(move |result| {
+                                let mut new_obj = result.clone();
+                                new_obj.insert(key.clone(), flattened_item.clone());
+                                new_obj
+                            })
+                        })
+                        .collect(),
+                    Value::Object(_) => flatten_json(val)
+                        .unwrap_or_default()
+                        .iter()
+                        .flat_map(|nested_result| {
+                            results.iter().map(move |result| {
+                                let mut new_obj = result.clone();
+                                new_obj.insert(key.clone(), nested_result.clone());
+                                new_obj
+                            })
                         })
-                    })
-                    .collect(),
-                Value::Object(_) => flatten_json(val)
-                    .iter()
-                    .flat_map(|nested_result| {
-                        results.iter().map(move |result| {
-                            let mut new_obj = result.clone();
-                            new_obj.insert(key.clone(), nested_result.clone());
-                            new_obj
+                        .collect(),
+                    _ => results
+                        .into_iter()
+                        .map(|mut result| {
+                            result.insert(key.clone(), val.clone());
+                            result
                         })
-                    })
-                    .collect(),
-                _ => results
-                    .into_iter()
-                    .map(|mut result| {
-                        result.insert(key.clone(), val.clone());
-                        result
-                    })
-                    .collect(),
-            })
-            .into_iter()
-            .map(Value::Object)
-            .collect(),
-        _ => vec![value.clone()],
+                        .collect(),
+                });
+
+            Ok(results.into_iter().map(Value::Object).collect())
+        }
+        _ => Ok(vec![value.clone()]),
+    }
+}
+
+/// recursively checks the level of nesting for the serde Value
+/// if Value has more than 4 levels of hierarchy, returns true
+/// example -
+/// 1. `{"a":{"b":{"c":{"d":{"e":["a","b"]}}}}}` ~> returns true
+/// 2. `{"a": [{"b": 1}, {"c": 2}], "d": {"e": 4}}` ~> returns false
+fn has_more_than_four_levels(value: &Value, current_level: usize) -> bool {
+    if current_level > 4 {
+        return true;
+    }
+    match value {
+        Value::Array(arr) => arr
+            .iter()
+            .any(|item| has_more_than_four_levels(item, current_level)),
+        Value::Object(map) => map
+            .values()
+            .any(|val| has_more_than_four_levels(val, current_level + 1)),
+        _ => false,
     }
 }
 
@@ -340,7 +372,9 @@ pub fn convert_to_array(flattened: Vec<Value>) -> Result<Value, JsonFlattenError
 
 #[cfg(test)]
 mod tests {
-    use crate::utils::json::flatten::flatten_array_objects;
+    use crate::utils::json::flatten::{
+        flatten_array_objects, flatten_json, has_more_than_four_levels,
+    };
 
     use super::{flatten, JsonFlattenError};
     use serde_json::{json, Map, Value};
@@ -598,4 +632,29 @@ mod tests {
             JsonFlattenError::FieldContainsPeriod(_)
         );
     }
+
+    #[test]
+    fn unacceptable_levels_of_nested_json() {
+        let value = json!({"a":{"b":{"c":{"d":{"e":["a","b"]}}}}});
+        assert_eq!(has_more_than_four_levels(&value, 1), true);
+    }
+
+    #[test]
+    fn acceptable_levels_of_nested_json() {
+        let value = json!({"a":{"b":{"e":["a","b"]}}});
+        assert_eq!(has_more_than_four_levels(&value, 1), false);
+    }
+
+    #[test]
+    fn flatten_json_success() {
+        let value = json!({"a":{"b":{"e":["a","b"]}}});
+        let expected = vec![json!({"a":{"b":{"e":"a"}}}), json!({"a":{"b":{"e":"b"}}})];
+        assert_eq!(flatten_json(&value).unwrap(), expected);
+    }
+
+    #[test]
+    fn flatten_json_error() {
+        let value = json!({"a":{"b":{"c":{"d":{"e":["a","b"]}}}}});
+        assert!(flatten_json(&value).is_err());
+    }
 }

src/utils/json/mod.rs

@@ -23,15 +23,20 @@ use serde_json::Value;
 
 pub mod flatten;
 
+/// calls the function `flatten_json` which results Vec<Value> or Error
+/// in case when Vec<Value> is returned, converts the Vec<Value> to Value of Array
+/// this is to ensure recursive flattening does not happen for heavily nested jsons
 pub fn flatten_json_body(
     body: &Value,
     time_partition: Option<&String>,
     time_partition_limit: Option<NonZeroU32>,
     custom_partition: Option<&String>,
     validation_required: bool,
 ) -> Result<Value, anyhow::Error> {
-    let mut nested_value = flatten::convert_to_array(flatten::flatten_json(body))?;
-
+    let mut nested_value = body.clone();
+    if let Ok(flattened_json) = flatten::flatten_json(body) {
+        nested_value = flatten::convert_to_array(flattened_json)?;
+    }
     flatten::flatten(
         &mut nested_value,
         "_",
@@ -40,7 +45,6 @@ pub fn flatten_json_body(
         custom_partition,
         validation_required,
     )?;
-
     Ok(nested_value)
 }
 
@@ -84,3 +88,29 @@ pub fn convert_to_string(value: &Value) -> Value {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::flatten_json_body;
+    use serde_json::json;
+
+    #[test]
+    fn hierarchical_json_flattening_success() {
+        let value = json!({"a":{"b":{"e":["a","b"]}}});
+        let expected = json!([{"a_b_e": "a"}, {"a_b_e": "b"}]);
+        assert_eq!(
+            flatten_json_body(&value, None, None, None, false).unwrap(),
+            expected
+        );
+    }
+
+    #[test]
+    fn hierarchical_json_flattening_failure() {
+        let value = json!({"a":{"b":{"c":{"d":{"e":["a","b"]}}}}});
+        let expected = json!({"a_b_c_d_e": ["a","b"]});
+        assert_eq!(
+            flatten_json_body(&value, None, None, None, false).unwrap(),
+            expected
+        );
+    }
+}