fix: flattening for kinesis

src/handlers/http/kinesis.rs

@@ -21,6 +21,8 @@ use serde::{Deserialize, Serialize};
 use serde_json::{Map, Value};
 use std::str;
 
+use crate::utils::json::flatten::{generic_flattening, has_more_than_max_allowed_levels};
+
 #[derive(Serialize, Deserialize, Debug)]
 #[serde(rename_all = "camelCase")]
 pub struct Message {
@@ -57,29 +59,35 @@ struct Data {
 //     "requestId": "b858288a-f5d8-4181-a746-3f3dd716be8a",
 //     "timestamp": "1704964113659"
 // }
-pub fn flatten_kinesis_logs(message: Message) -> Vec<Value> {
+pub async fn flatten_kinesis_logs(message: Message) -> Result<Vec<Value>, anyhow::Error> {
     let mut vec_kinesis_json = Vec::new();
 
     for record in message.records.iter() {
         let bytes = STANDARD.decode(record.data.clone()).unwrap();
-        let json_string: String = String::from_utf8(bytes).unwrap();
-        let json: serde_json::Value = serde_json::from_str(&json_string).unwrap();
-        let mut kinesis_json: Map<String, Value> = match serde_json::from_value(json) {
-            Ok(value) => value,
-            Err(error) => panic!("Failed to deserialize JSON: {}", error),
-        };
-
-        kinesis_json.insert(
-            "requestId".to_owned(),
-            Value::String(message.request_id.clone()),
-        );
-        kinesis_json.insert(
-            "timestamp".to_owned(),
-            Value::String(message.timestamp.to_string()),
-        );
+        if let Ok(json_string) = String::from_utf8(bytes) {
+            let json: serde_json::Value = serde_json::from_str(&json_string).unwrap();
+            if !has_more_than_max_allowed_levels(&json, 1) {
+                let flattened_json_arr = generic_flattening(&json)?;
+                for flattened_json in flattened_json_arr {
+                    let mut kinesis_json: Map<String, Value> =
+                        match serde_json::from_value(flattened_json) {
+                            Ok(value) => value,
+                            Err(error) => panic!("Failed to deserialize JSON: {}", error),
+                        };
+                    kinesis_json.insert(
+                        "requestId".to_owned(),
+                        Value::String(message.request_id.clone()),
+                    );
+                    kinesis_json.insert(
+                        "timestamp".to_owned(),
+                        Value::String(message.timestamp.to_string()),
+                    );
 
-        vec_kinesis_json.push(Value::Object(kinesis_json));
+                    vec_kinesis_json.push(Value::Object(kinesis_json));
+                }
+            }
+        }
     }
 
-    vec_kinesis_json
+    Ok(vec_kinesis_json)
 }

src/handlers/http/modal/utils/ingest_utils.rs

@@ -61,9 +61,9 @@ pub async fn flatten_and_push_logs(
         LogSource::Kinesis => {
             //custom flattening required for Amazon Kinesis
             let message: Message = serde_json::from_value(json)?;
-            for record in flatten_kinesis_logs(message) {
-                push_logs(stream_name, record, log_source, p_custom_fields).await?;
-            }
+            let flattened_kinesis_data = flatten_kinesis_logs(message).await?;
+            let record = convert_to_array(flattened_kinesis_data)?;
+            push_logs(stream_name, record, log_source, p_custom_fields).await?;
         }
         LogSource::OtelLogs => {
             //custom flattening required for otel logs

src/utils/json/mod.rs

@@ -44,10 +44,8 @@ pub fn flatten_json_body(
     // Flatten the json body only if new schema and has less than 4 levels of nesting
     let mut nested_value = if schema_version == SchemaVersion::V1
         && !has_more_than_max_allowed_levels(&body, 1)
-        && matches!(
-            log_source,
-            LogSource::Json | LogSource::Custom(_) | LogSource::Kinesis
-        ) {
+        && matches!(log_source, LogSource::Json | LogSource::Custom(_))
+    {
         let flattened_json = generic_flattening(&body)?;
         convert_to_array(flattened_json)?
     } else {