From acf09411ea5d530bffe93873f1c73ad5a1474bd8 Mon Sep 17 00:00:00 2001
From: juan-coralogix <89215136+juan-coralogix@users.noreply.github.com>
Date: Thu, 14 Mar 2024 11:51:43 -0300
Subject: [PATCH] Custom Lambda Addition (#74)

---
 CHANGELOG.md     |  3 +++
 README.md        |  2 +-
 src/coralogix.rs | 31 +++++++++++++++++++++++++++++--
 template.yaml    | 12 +++++++++++-
 4 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5dc2a7fe..3abadfe7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,4 +1,7 @@
 # Changelog
+## v1.0.3 / 
+### 🚀 New components 🚀
+- Custom Metadata can be added to the log messages.
 
 ## v1.0.2 / 2024-03-06
 - Update dependencies to fix security vulnerabilities
diff --git a/README.md b/README.md
index f8ed8dc1..45c5229d 100644
--- a/README.md
+++ b/README.md
@@ -172,7 +172,7 @@ These are optional parameters if you wish to receive notification emails, exclud
 | BlockingPattern   | Enter a regular expression to identify lines excluded from being sent to Coralogix. For example, use `MainActivity.java:\d{3}` to match log lines with `MainActivity` followed by exactly three digits. |               |                    |
 | SamplingRate      | Send messages at a specific rate, such as 1 out of every N logs. For example, if your value is 10, a message will be sent for every 10th log.                                                           | 1             | :heavy_check_mark: |
 | AddMetadata       | Add metadata to the log message. Expects comma separated values. Options for S3 are `bucket_name`,`key_name`. For CloudWatch use `stream_name`, `loggroup_name` .                                       |               |                    |
-
+| CustomMetadata    | Add custom metadata to the log message. Expects comma separated values. Options are key1=value1,key2=value2 | | |
 ### Lambda Configuration (Optional)
 
 These are the default presets for Lambda. Read [Troubleshooting](#troubleshooting) for more information on changing these defaults.
diff --git a/src/coralogix.rs b/src/coralogix.rs
index 14d7260a..dd869662 100644
--- a/src/coralogix.rs
+++ b/src/coralogix.rs
@@ -13,6 +13,7 @@ use std::time::Instant;
 use std::vec::Vec;
 use time::OffsetDateTime;
 use tracing::{error, info};
+use std::env;
 
 pub async fn process_batches(
     logs: Vec<String>,
@@ -110,6 +111,8 @@ struct JsonMessage {
     bucket_name: Option<String>,
     #[serde(skip_serializing_if = "Option::is_none")]
     key_name: Option<String>,
+    #[serde(flatten)]
+    custom_metadata: HashMap<String, String>,
 }
 
 fn convert_to_log_entry(
@@ -148,6 +151,7 @@ fn convert_to_log_entry(
         loggroup_name: None,
         bucket_name: None,
         key_name: None,
+        custom_metadata: HashMap::new(),
     };
 
     let add_metadata: Vec<&str> = config.add_metadata.split(',').map(|s| s.trim()).collect();
@@ -178,8 +182,31 @@ fn convert_to_log_entry(
             }
         }
     }
-
-    let body =  if message.stream_name.is_some() || message.loggroup_name.is_some() || message.bucket_name.is_some() || message.key_name.is_some() {
+    if let Ok(custom_metadata_str) = env::var("CUSTOM_METADATA") {
+        debug!("Custom metadata STR: {}", custom_metadata_str);
+        let mut metadata = HashMap::new();
+        let pairs = custom_metadata_str.split(',');
+    
+        for pair in pairs {
+            let split_pair: Vec<&str> = pair.split('=').collect();
+            match split_pair.as_slice() {
+                [key, value] => {
+                    metadata.insert(key.to_string(), value.to_string());
+                },
+                _ => {
+                    error!("Failed to split key-value pair: {}", pair);
+                    continue;
+                }
+            }
+        }
+    
+        if !metadata.is_empty() {
+            debug!("Custom metadata: {:?}", metadata);
+            message.custom_metadata = metadata;
+        }
+    }
+    debug!("Message metadata: {:?}", message.custom_metadata);
+    let body =  if message.stream_name.is_some() || message.loggroup_name.is_some() || message.bucket_name.is_some() || message.key_name.is_some() || !message.custom_metadata.is_empty() {
         serde_json::to_value(&message).unwrap_or(message.message)
     } else {
         message.message
diff --git a/template.yaml b/template.yaml
index b2fbc338..123da561 100644
--- a/template.yaml
+++ b/template.yaml
@@ -82,6 +82,7 @@ Metadata:
           - BlockingPattern
           - SamplingRate
           - AddMetadata
+          - CustomMetadata
       - Label:
           default: Lambda configuration
         Parameters:
@@ -145,7 +146,11 @@ Parameters:
     Description: |
       Add metadata to the log message. Expects comma separated values. Options are bucket_name,key_name,stream_name
     Default: ''
-
+  CustomMetadata:
+    Type: String
+    Description: |
+      Add custom metadata to the log message. Expects comma separated values. Options are key1=value1,key2=value2
+    Default: ''
   BlockingPattern:
     Type: String
     Description: Regular expression to detect lines that should be excluded from sent to Coralogix
@@ -358,6 +363,7 @@ Conditions:
   CSVDelimiterUse: !Equals [!Ref IntegrationType, 'S3Csv']
   NewlinePatternNotSet: !Equals [ !Ref NewlinePattern, '' ]
   AddMetadataNotSet: !Equals [ !Ref AddMetadata, '']
+  CustomMetadataNotSet: !Equals [ !Ref CustomMetadata, '']
   IsSNSIntegration: !Equals [ !Ref IntegrationType, 'Sns' ]
   UseECRScan: !Equals [ !Ref IntegrationType, 'EcrScan' ]
   IsSQSIntegration: !Equals [ !Ref IntegrationType, 'Sqs' ]
@@ -570,6 +576,10 @@ Globals:
           - AddMetadataNotSet
           - !Ref AWS::NoValue
           - !Ref AddMetadata
+        CUSTOM_METADATA: !If
+          - CustomMetadataNotSet
+          - !Ref AWS::NoValue
+          - !Ref CustomMetadata
         NEWLINE_PATTERN: !If
           - NewlinePatternNotSet
           - !Ref AWS::NoValue