fix for #605: added feature for flattening kinesis logs before ingesting

Author: nikhilsinhaparseable

server/src/handlers.rs

@@ -24,6 +24,7 @@ const PREFIX_META: &str = "x-p-meta-";
 const STREAM_NAME_HEADER_KEY: &str = "x-p-stream";
 const AUTHORIZATION_KEY: &str = "authorization";
 const KINESIS_COMMON_ATTRIBUTES_KEY: &str = "x-amz-firehose-common-attributes";
+const LOG_SOURCE_KEY: &str = "x-p-log-source";
 const SEPARATOR: char = '^';
 
 const OIDC_SCOPE: &str = "openid profile email";

server/src/handlers/http.rs

@@ -47,6 +47,7 @@ mod oidc;
 mod query;
 mod rbac;
 mod role;
+mod kinesis;
 
 include!(concat!(env!("OUT_DIR"), "/generated.rs"));
 

server/src/handlers/http/ingest.rs

@@ -29,11 +29,12 @@ use serde_json::Value;
 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};
+use crate::handlers::{PREFIX_META, PREFIX_TAGS, SEPARATOR, STREAM_NAME_HEADER_KEY, LOG_SOURCE_KEY};
 use crate::metadata::STREAM_INFO;
 use crate::utils::header_parsing::{collect_labelled_headers, ParseHeaderError};
 
 use super::logstream::error::CreateStreamError;
+use super::kinesis;
 
 // Handler for POST /api/v1/ingest
 // ingests events by extracting stream name from header
@@ -46,13 +47,35 @@ 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?;
         Ok(HttpResponse::Ok().finish())
     } else {
         Err(PostError::Header(ParseHeaderError::MissingStreamName))
     }
 }
 
+fn flatten_logs(req: HttpRequest, body: &Bytes){
+    //flatten logs
+    if let Some((_, log_source)) = req
+        .headers()
+        .iter()
+        .find(|&(key, _)| key == LOG_SOURCE_KEY)
+    {
+        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
+        }; 
+    }
+}
+
+
 // Handler for POST /api/v1/logstream/{logstream}
 // only ingests events into the specified logstream
 // fails if the logstream does not exist
@@ -70,6 +93,7 @@ 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)?
     };
 
@@ -94,12 +118,15 @@ 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

@@ -0,0 +1,29 @@
+use bytes::Bytes;
+use serde::{Deserialize, Serialize};
+use base64::{engine::general_purpose::STANDARD, Engine as _};
+use std::str;
+
+
+#[derive(Serialize, Deserialize, Debug)]
+struct Message {
+    #[serde(rename = "records_data")]
+    records_data: Vec<String>,
+    #[serde(rename = "requestId")]
+    request_id: String,
+    timestamp: u64
+}
+
+pub fn flatten_kinesis_logs(body: &Bytes){
+    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);
+
+
+    for record in message.records_data.iter() {
+        let bytes = STANDARD.decode(record).unwrap();
+        println!("{:?}", String::from_utf8(bytes).unwrap());
+
+    }
+}