fix for #605: Decode and flatten logs from kinesis before ingesting on Parseable

Author: nikhilsinhaparseable

server/src/handlers.rs

@@ -31,3 +31,7 @@ const OIDC_SCOPE: &str = "openid profile email";
 const COOKIE_AGE_DAYS: usize = 7;
 const SESSION_COOKIE_NAME: &str = "session";
 const USER_COOKIE_NAME: &str = "username";
+
+// constants for Log Source values for known sources
+const LOG_SOURCE_VALUE_FOR_KINEIS: &str = "kinesis";
+const LOG_SOURCE_VALUE_FOR_OTEL: &str = "otel";

server/src/handlers/http.rs

@@ -40,14 +40,14 @@ use self::middleware::{DisAllowRootUser, RouteExt};
 mod about;
 mod health_check;
 mod ingest;
+mod kinesis;
 mod llm;
 mod logstream;
 mod middleware;
 mod oidc;
 mod query;
 mod rbac;
 mod role;
-mod kinesis;
 
 include!(concat!(env!("OUT_DIR"), "/generated.rs"));
 

server/src/handlers/http/ingest.rs

@@ -16,25 +16,26 @@
  *
  */
 
-use std::collections::HashMap;
-use std::sync::Arc;
-
-use actix_web::http::header::ContentType;
-use actix_web::{HttpRequest, HttpResponse};
+use actix_web::{http::header::ContentType, HttpRequest, HttpResponse};
 use arrow_schema::Field;
 use bytes::Bytes;
 use http::StatusCode;
 use serde_json::Value;
+use std::collections::{BTreeMap,HashMap};
+use std::sync::Arc;
 
 use crate::event::error::EventError;
 use crate::event::format::EventFormat;
 use crate::event::{self, format};
-use crate::handlers::{PREFIX_META, PREFIX_TAGS, SEPARATOR, STREAM_NAME_HEADER_KEY, LOG_SOURCE_KEY};
+use crate::handlers::{
+    LOG_SOURCE_KEY, PREFIX_META, PREFIX_TAGS, SEPARATOR, STREAM_NAME_HEADER_KEY, 
+    LOG_SOURCE_VALUE_FOR_KINEIS, LOG_SOURCE_VALUE_FOR_OTEL
+};
 use crate::metadata::STREAM_INFO;
 use crate::utils::header_parsing::{collect_labelled_headers, ParseHeaderError};
 
-use super::logstream::error::CreateStreamError;
 use super::kinesis;
+use super::logstream::error::CreateStreamError;
 
 // Handler for POST /api/v1/ingest
 // ingests events by extracting stream name from header
@@ -48,40 +49,44 @@ pub async fn ingest(req: HttpRequest, body: Bytes) -> Result<HttpResponse, PostE
         let stream_name = stream_name.to_str().unwrap().to_owned();
         create_stream_if_not_exists(&stream_name).await?;
 
-        //section to flatten ingested log data
-        let cloned_req = req.clone();
-        flatten_logs(cloned_req, &body);
-        //section ends
-        push_logs(stream_name, req, body).await?;
+        flatten_and_push_logs(req, body, stream_name).await?;
         Ok(HttpResponse::Ok().finish())
     } else {
         Err(PostError::Header(ParseHeaderError::MissingStreamName))
     }
 }
 
-fn flatten_logs(req: HttpRequest, body: &Bytes){
+async fn flatten_and_push_logs(
+    req: HttpRequest,
+    body: Bytes,
+    stream_name: String,
+) -> Result<(), PostError> {
     //flatten logs
-    if let Some((_, log_source)) = req
-        .headers()
-        .iter()
-        .find(|&(key, _)| key == LOG_SOURCE_KEY)
-    {
+    if let Some((_, log_source)) = req.headers().iter().find(|&(key, _)| key == LOG_SOURCE_KEY) {
+        let mut json: Vec<BTreeMap<String, Value>> = Vec::new();
         let log_source: String = log_source.to_str().unwrap().to_owned();
-        //println!("log source: {}", log_source);
-        match log_source.as_str(){ 
-            "kinesis_firehose" => kinesis::flatten_kinesis_logs(body),
-            _ => {}//do nothing so far
-        }; 
+        match log_source.as_str() {
+            LOG_SOURCE_VALUE_FOR_KINEIS => json = kinesis::flatten_kinesis_logs(&body),
+            LOG_SOURCE_VALUE_FOR_OTEL => {}
+            _ => {}
+        }
+        for record in json.iter_mut() {
+            let body: Bytes = serde_json::to_vec(record).unwrap().into();
+            push_logs(stream_name.to_string(), req.clone(), body).await?;
+        }
+    } else {
+        push_logs(stream_name.to_string(), req, body).await?;
     }
+    Ok(())
 }
 
-
 // Handler for POST /api/v1/logstream/{logstream}
 // only ingests events into the specified logstream
 // fails if the logstream does not exist
 pub async fn post_event(req: HttpRequest, body: Bytes) -> Result<HttpResponse, PostError> {
     let stream_name: String = req.match_info().get("logstream").unwrap().parse().unwrap();
-    push_logs(stream_name, req, body).await?;
+
+    flatten_and_push_logs(req, body, stream_name).await?;
     Ok(HttpResponse::Ok().finish())
 }
 
@@ -93,7 +98,6 @@ async fn push_logs(stream_name: String, req: HttpRequest, body: Bytes) -> Result
             .ok_or(PostError::StreamNotFound(stream_name.clone()))?
             .schema
             .clone();
-        
         into_event_batch(req, body, schema)?
     };
 
@@ -118,15 +122,12 @@ fn into_event_batch(
     let tags = collect_labelled_headers(&req, PREFIX_TAGS, SEPARATOR)?;
     let metadata = collect_labelled_headers(&req, PREFIX_META, SEPARATOR)?;
     let size = body.len();
-    println!("{:?}", body);
     let body: Value = serde_json::from_slice(&body)?;
     let event = format::json::Event {
         data: body,
         tags,
         metadata,
     };
-    
-    println!("{:?}", event);
     let (rb, is_first) = event.into_recordbatch(schema)?;
     Ok((size, rb, is_first))
 }

server/src/handlers/http/kinesis.rs

@@ -1,29 +1,48 @@
+use base64::{engine::general_purpose::STANDARD, Engine as _};
 use bytes::Bytes;
 use serde::{Deserialize, Serialize};
-use base64::{engine::general_purpose::STANDARD, Engine as _};
+use serde_json::Value;
+use std::collections::BTreeMap;
 use std::str;
 
-
 #[derive(Serialize, Deserialize, Debug)]
 struct Message {
-    #[serde(rename = "records_data")]
-    records_data: Vec<String>,
+    #[serde(rename = "records")]
+    records: Vec<Data>,
     #[serde(rename = "requestId")]
     request_id: String,
-    timestamp: u64
+    timestamp: u64,
+}
+#[derive(Serialize, Deserialize, Debug)]
+struct Data {
+    #[serde(rename = "data")]
+    data: String,
 }
 
-pub fn flatten_kinesis_logs(body: &Bytes){
+pub fn flatten_kinesis_logs(body: &Bytes) -> Vec<BTreeMap<String, Value>> {
     let body_str = std::str::from_utf8(body).unwrap();
     let message: Message = serde_json::from_str(body_str).unwrap();
-    println!("{:?}", message.records_data);
-    println!("{:?}", message.request_id);
-    println!("{:?}", message.timestamp);
+    let mut vec_kinesis_json: Vec<BTreeMap<String, Value>> = 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: BTreeMap<String, Value> = match serde_json::from_value(json) {
+            Ok(value) => value,
+            Err(error) => panic!("Failed to deserialize JSON: {}", error),
+        };
 
-    for record in message.records_data.iter() {
-        let bytes = STANDARD.decode(record).unwrap();
-        println!("{:?}", String::from_utf8(bytes).unwrap());
+        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(kinesis_json);
     }
-}
+    vec_kinesis_json
+}

server/src/handlers/http/middleware.rs

@@ -27,7 +27,10 @@ use actix_web::{
 };
 use futures_util::future::LocalBoxFuture;
 
-use crate::handlers::{AUTHORIZATION_KEY, KINESIS_COMMON_ATTRIBUTES_KEY, STREAM_NAME_HEADER_KEY};
+use crate::handlers::{
+    AUTHORIZATION_KEY, KINESIS_COMMON_ATTRIBUTES_KEY, LOG_SOURCE_KEY, 
+    STREAM_NAME_HEADER_KEY, LOG_SOURCE_VALUE_FOR_KINEIS
+};
 use crate::{
     option::CONFIG,
     rbac::Users,
@@ -149,6 +152,10 @@ where
                 header::HeaderValue::from_str(&message.common_attributes.x_p_stream.clone())
                     .unwrap(),
             );
+            req.headers_mut().insert(
+                HeaderName::from_static(LOG_SOURCE_KEY),
+                header::HeaderValue::from_static(LOG_SOURCE_VALUE_FOR_KINEIS),
+            );
         }
 
         /* ## Section end */