suggestions

Author: Devdutt Shenoi

Cargo.toml

@@ -64,9 +64,10 @@ humantime-serde = "1.1"
 itertools = "0.13.0"
 num_cpus = "1.15"
 once_cell = "1.17.1"
+opentelemetry-proto = "0.27.0"
 prometheus = { version = "0.13", features = ["process"] }
 rand = "0.8.5"
-rdkafka = {version = "0.36.2", default-features = false, features = ["tokio"]}
+rdkafka = { version = "0.36.2", default-features = false, features = ["tokio"] }
 regex = "1.7.3"
 relative-path = { version = "1.7", features = ["serde"] }
 reqwest = { version = "0.11.27", default-features = false, features = [
@@ -106,7 +107,6 @@ prost = "0.13.3"
 prometheus-parse = "0.2.5"
 sha2 = "0.10.8"
 tracing = "0.1.41"
-opentelemetry-proto = "0.27.0"
 
 [build-dependencies]
 cargo_toml = "0.20.1"

src/handlers/http/ingest.rs

@@ -41,6 +41,10 @@ use arrow_schema::Schema;
 use bytes::Bytes;
 use chrono::Utc;
 use http::StatusCode;
+use nom::AsBytes;
+use opentelemetry_proto::tonic::logs::v1::LogsData;
+use opentelemetry_proto::tonic::metrics::v1::MetricsData;
+use opentelemetry_proto::tonic::trace::v1::TracesData;
 use serde_json::Value;
 use std::collections::HashMap;
 use std::sync::Arc;
@@ -113,23 +117,20 @@ pub async fn handle_otel_logs_ingestion(
     req: HttpRequest,
     body: Bytes,
 ) -> Result<HttpResponse, PostError> {
-    if let Some((_, stream_name)) = req
-        .headers()
-        .iter()
-        .find(|&(key, _)| key == STREAM_NAME_HEADER_KEY)
-    {
-        let stream_name = stream_name.to_str().unwrap().to_owned();
-        create_stream_if_not_exists(&stream_name, &StreamType::UserDefined.to_string()).await?;
-
-        //custom flattening required for otel logs
-        let mut json = flatten_otel_logs(&body);
-        for record in json.iter_mut() {
-            let body: Bytes = serde_json::to_vec(record).unwrap().into();
-            push_logs(&stream_name, &req, &body).await?;
-        }
-    } else {
+    let Some(stream_name) = req.headers().get(STREAM_NAME_HEADER_KEY) else {
         return Err(PostError::Header(ParseHeaderError::MissingStreamName));
+    };
+    let stream_name = stream_name.to_str().unwrap().to_owned();
+    create_stream_if_not_exists(&stream_name, &StreamType::UserDefined.to_string()).await?;
+
+    //custom flattening required for otel logs
+    let logs: LogsData = serde_json::from_slice(body.as_bytes())?;
+    let mut json = flatten_otel_logs(&logs);
+    for record in json.iter_mut() {
+        let body: Bytes = serde_json::to_vec(record).unwrap().into();
+        push_logs(&stream_name, &req, &body).await?;
     }
+
     Ok(HttpResponse::Ok().finish())
 }
 
@@ -140,23 +141,20 @@ pub async fn handle_otel_metrics_ingestion(
     req: HttpRequest,
     body: Bytes,
 ) -> Result<HttpResponse, PostError> {
-    if let Some((_, stream_name)) = req
-        .headers()
-        .iter()
-        .find(|&(key, _)| key == STREAM_NAME_HEADER_KEY)
-    {
-        let stream_name = stream_name.to_str().unwrap().to_owned();
-        create_stream_if_not_exists(&stream_name, &StreamType::UserDefined.to_string()).await?;
-
-        //custom flattening required for otel metrics
-        let mut json = flatten_otel_metrics(&body);
-        for record in json.iter_mut() {
-            let body: Bytes = serde_json::to_vec(record).unwrap().into();
-            push_logs(&stream_name, &req, &body).await?;
-        }
-    } else {
+    let Some(stream_name) = req.headers().get(STREAM_NAME_HEADER_KEY) else {
         return Err(PostError::Header(ParseHeaderError::MissingStreamName));
+    };
+    let stream_name = stream_name.to_str().unwrap().to_owned();
+    create_stream_if_not_exists(&stream_name, &StreamType::UserDefined.to_string()).await?;
+
+    //custom flattening required for otel metrics
+    let metrics: MetricsData = serde_json::from_slice(body.as_bytes())?;
+    let mut json = flatten_otel_metrics(metrics);
+    for record in json.iter_mut() {
+        let body: Bytes = serde_json::to_vec(record).unwrap().into();
+        push_logs(&stream_name, &req, &body).await?;
     }
+
     Ok(HttpResponse::Ok().finish())
 }
 
@@ -167,23 +165,20 @@ pub async fn handle_otel_traces_ingestion(
     req: HttpRequest,
     body: Bytes,
 ) -> Result<HttpResponse, PostError> {
-    if let Some((_, stream_name)) = req
-        .headers()
-        .iter()
-        .find(|&(key, _)| key == STREAM_NAME_HEADER_KEY)
-    {
-        let stream_name = stream_name.to_str().unwrap().to_owned();
-        create_stream_if_not_exists(&stream_name, &StreamType::UserDefined.to_string()).await?;
-
-        //custom flattening required for otel traces
-        let mut json = flatten_otel_traces(&body);
-        for record in json.iter_mut() {
-            let body: Bytes = serde_json::to_vec(record).unwrap().into();
-            push_logs(&stream_name, &req, &body).await?;
-        }
-    } else {
+    let Some(stream_name) = req.headers().get(STREAM_NAME_HEADER_KEY) else {
         return Err(PostError::Header(ParseHeaderError::MissingStreamName));
+    };
+    let stream_name = stream_name.to_str().unwrap().to_owned();
+    create_stream_if_not_exists(&stream_name, &StreamType::UserDefined.to_string()).await?;
+
+    //custom flattening required for otel traces
+    let traces: TracesData = serde_json::from_slice(body.as_bytes())?;
+    let mut json = flatten_otel_traces(&traces);
+    for record in json.iter_mut() {
+        let body: Bytes = serde_json::to_vec(record).unwrap().into();
+        push_logs(&stream_name, &req, &body).await?;
     }
+
     Ok(HttpResponse::Ok().finish())
 }
 

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

@@ -21,6 +21,10 @@ use anyhow::anyhow;
 use arrow_schema::Field;
 use bytes::Bytes;
 use chrono::{DateTime, NaiveDateTime, Utc};
+use nom::AsBytes;
+use opentelemetry_proto::tonic::{
+    logs::v1::LogsData, metrics::v1::MetricsData, trace::v1::TracesData,
+};
 use itertools::Itertools;
 use serde_json::Value;
 use std::{
@@ -49,39 +53,39 @@ pub async fn flatten_and_push_logs(
     body: Bytes,
     stream_name: &str,
 ) -> Result<(), PostError> {
-    //flatten logs
-    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();
-        match log_source.as_str() {
-            LOG_SOURCE_KINESIS => json = kinesis::flatten_kinesis_logs(&body),
-
-            //custom flattening required for otel logs
-            LOG_SOURCE_OTEL_LOGS => {
-                json = flatten_otel_logs(&body);
-            }
-
-            //custom flattening required for otel metrics
-            LOG_SOURCE_OTEL_METRICS => {
-                json = flatten_otel_metrics(&body);
-            }
-
-            //custom flattening required for otel traces
-            LOG_SOURCE_OTEL_TRACES => {
-                json = flatten_otel_traces(&body);
-            }
-            _ => {
-                tracing::warn!("Unknown log source: {}", log_source);
-                push_logs(stream_name, &req, &body).await?;
-            }
+    let Some(log_source) = req.headers().get(LOG_SOURCE_KEY) else {
+        push_logs(stream_name, &req, &body).await?;
+        return Ok(());
+    };
+    let mut json: Vec<BTreeMap<String, Value>> = Vec::new();
+    match log_source.to_str().unwrap() {
+        LOG_SOURCE_KINESIS => json = kinesis::flatten_kinesis_logs(&body),
+        //custom flattening required for otel logs
+        LOG_SOURCE_OTEL_LOGS => {
+            let logs: LogsData = serde_json::from_slice(body.as_bytes())?;
+            json = flatten_otel_logs(&logs);
+        }
+        //custom flattening required for otel metrics
+        LOG_SOURCE_OTEL_METRICS => {
+            let metrics: MetricsData = serde_json::from_slice(body.as_bytes())?;
+            json = flatten_otel_metrics(metrics);
+        }
+        //custom flattening required for otel traces
+        LOG_SOURCE_OTEL_TRACES => {
+            let traces: TracesData = serde_json::from_slice(body.as_bytes())?;
+            json = flatten_otel_traces(&traces);
         }
-        for record in json.iter_mut() {
-            let body: Bytes = serde_json::to_vec(record).unwrap().into();
+        log_source => {
+            tracing::warn!("Unknown log source: {}", log_source);
             push_logs(stream_name, &req, &body).await?;
         }
-    } else {
+    }
+
+    for record in json.iter_mut() {
+        let body: Bytes = serde_json::to_vec(record).unwrap().into();
         push_logs(stream_name, &req, &body).await?;
     }
+
     Ok(())
 }
 

src/otel/logs.rs

@@ -16,7 +16,6 @@
  *
  */
 
-use bytes::Bytes;
 use opentelemetry_proto::tonic::logs::v1::LogRecord;
 use opentelemetry_proto::tonic::logs::v1::LogsData;
 use opentelemetry_proto::tonic::logs::v1::ScopeLogs;
@@ -125,11 +124,8 @@ fn flatten_scope_log(scope_log: &ScopeLogs) -> Vec<BTreeMap<String, Value>> {
 
 /// this function performs the custom flattening of the otel logs
 /// and returns a `Vec` of `BTreeMap` of the flattened json
-pub fn flatten_otel_logs(body: &Bytes) -> Vec<BTreeMap<String, Value>> {
-    let body_str = std::str::from_utf8(body).unwrap();
-    let message: LogsData = serde_json::from_str(body_str).unwrap();
+pub fn flatten_otel_logs(message: &LogsData) -> Vec<BTreeMap<String, Value>> {
     let mut vec_otel_json = Vec::new();
-
     for record in &message.resource_logs {
         let mut resource_log_json = BTreeMap::new();
 

src/otel/metrics.rs

@@ -18,7 +18,6 @@
 
 use std::collections::BTreeMap;
 
-use bytes::Bytes;
 use opentelemetry_proto::tonic::metrics::v1::number_data_point::Value as NumberDataPointValue;
 use opentelemetry_proto::tonic::metrics::v1::{
     exemplar::Value as ExemplarValue, exponential_histogram_data_point::Buckets, metric, Exemplar,
@@ -386,9 +385,7 @@ pub fn flatten_metrics_record(metrics_record: &Metric) -> Vec<BTreeMap<String, V
 
 /// this function performs the custom flattening of the otel metrics
 /// and returns a `Vec` of `BTreeMap` of the flattened json
-pub fn flatten_otel_metrics(body: &Bytes) -> Vec<BTreeMap<String, Value>> {
-    let body_str = std::str::from_utf8(body).unwrap();
-    let message: MetricsData = serde_json::from_str(body_str).unwrap();
+pub fn flatten_otel_metrics(message: MetricsData) -> Vec<BTreeMap<String, Value>> {
     let mut vec_otel_json = Vec::new();
     for record in &message.resource_metrics {
         let mut resource_metrics_json = BTreeMap::new();

src/otel/traces.rs

@@ -16,8 +16,6 @@
  *
  */
 
-use bytes::Bytes;
-
 use opentelemetry_proto::tonic::trace::v1::span::Event;
 use opentelemetry_proto::tonic::trace::v1::span::Link;
 use opentelemetry_proto::tonic::trace::v1::ScopeSpans;
@@ -71,9 +69,7 @@ fn flatten_scope_span(scope_span: &ScopeSpans) -> Vec<BTreeMap<String, Value>> {
 
 /// this function performs the custom flattening of the otel traces event
 /// and returns a `Vec` of `BTreeMap` of the flattened json
-pub fn flatten_otel_traces(body: &Bytes) -> Vec<BTreeMap<String, Value>> {
-    let body_str = std::str::from_utf8(body).unwrap();
-    let message: TracesData = serde_json::from_str(body_str).unwrap();
+pub fn flatten_otel_traces(message: &TracesData) -> Vec<BTreeMap<String, Value>> {
     let mut vec_otel_json = Vec::new();
 
     for record in &message.resource_spans {