From 7f6aef9ecba061ea28352b8a0bbdd7ac9ece64e2 Mon Sep 17 00:00:00 2001
From: Ashish Kolhe <35160958+oasisk@users.noreply.github.com>
Date: Fri, 7 Apr 2023 17:33:54 +0530
Subject: [PATCH] implement enum (#519)

---
 src/meta/ingestion.rs                |  7 +++++++
 src/service/logs/kinesis_firehose.rs | 13 ++++++++-----
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/src/meta/ingestion.rs b/src/meta/ingestion.rs
index c4dabe420b..509dd5a0a3 100644
--- a/src/meta/ingestion.rs
+++ b/src/meta/ingestion.rs
@@ -213,3 +213,10 @@ pub struct KinesisFHIngestionResponse {
     pub error_message: Option<String>,
     pub timestamp: i64,
 }
+
+#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
+#[serde(rename_all = "lowercase")]
+pub enum AWSRecordType {
+    JSON,
+    Cloudwatch,
+}
diff --git a/src/service/logs/kinesis_firehose.rs b/src/service/logs/kinesis_firehose.rs
index 3ca80578cd..e120699317 100644
--- a/src/service/logs/kinesis_firehose.rs
+++ b/src/service/logs/kinesis_firehose.rs
@@ -20,6 +20,7 @@ use crate::infra::file_lock;
 use crate::infra::metrics;
 use crate::meta::alert::{Alert, Trigger};
 use crate::meta::http::HttpResponse as MetaHttpResponse;
+use crate::meta::ingestion::AWSRecordType;
 use crate::meta::ingestion::KinesisFHData;
 use crate::meta::ingestion::KinesisFHIngestionResponse;
 use crate::meta::ingestion::KinesisFHRequest;
@@ -113,10 +114,10 @@ pub async fn process(
     let mut buf: AHashMap<String, Vec<String>> = AHashMap::new();
     for record in request.records {
         match decode_and_decompress(&record.data) {
-            Ok((decompressed_data, is_json)) => {
+            Ok((decompressed_data, record_type)) => {
                 let mut value = json::Value::Null;
                 let mut timestamp = 0;
-                if !is_json {
+                if record_type.eq(&AWSRecordType::Cloudwatch) {
                     let kfh_data: KinesisFHData = json::from_str(&decompressed_data)?;
 
                     for event in kfh_data.log_events.iter() {
@@ -349,13 +350,15 @@ pub async fn process(
     }))
 }
 
-fn decode_and_decompress(encoded_data: &str) -> Result<(String, bool), Box<dyn std::error::Error>> {
+fn decode_and_decompress(
+    encoded_data: &str,
+) -> Result<(String, AWSRecordType), Box<dyn std::error::Error>> {
     let decoded_data = crate::common::base64::decode_raw(encoded_data)?;
     let mut gz = GzDecoder::new(&decoded_data[..]);
     let mut decompressed_data = String::new();
     match gz.read_to_string(&mut decompressed_data) {
-        Ok(_) => Ok((decompressed_data, false)),
-        Err(_) => Ok((String::from_utf8(decoded_data)?, true)),
+        Ok(_) => Ok((decompressed_data, AWSRecordType::Cloudwatch)),
+        Err(_) => Ok((String::from_utf8(decoded_data)?, AWSRecordType::JSON)),
     }
 }