From 2b31f9500948da334ab9f613155d856da7249883 Mon Sep 17 00:00:00 2001
From: "laszlo.bica" <lbica@cmtelematics.com>
Date: Tue, 6 Sep 2022 17:07:52 +0200
Subject: [PATCH] Address feature request in Issue 13
 https://github.com/SumoLogic/sumologic-lambda-extensions/issues/13

Implement a feature switch to turn off log enhancement for JSON logs
Feature switch defaults to `true`, so behaviour is unchanged
---
 lambda-extensions/config/config.go         | 16 ++++++++++++++--
 lambda-extensions/sumoclient/sumoclient.go |  8 ++++++--
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/lambda-extensions/config/config.go b/lambda-extensions/config/config.go
index a6c8bea..53fc596 100644
--- a/lambda-extensions/config/config.go
+++ b/lambda-extensions/config/config.go
@@ -34,6 +34,7 @@ type LambdaExtensionConfig struct {
 	MaxDataPayloadSize     int
 	LambdaRegion           string
 	SourceCategoryOverride string
+	EnhanceJsonLogs        bool
 }
 
 var defaultLogTypes = []string{"platform", "function"}
@@ -73,6 +74,7 @@ func (cfg *LambdaExtensionConfig) setDefaults() {
 	maxConcurrentRequests := os.Getenv("SUMO_MAX_CONCURRENT_REQUESTS")
 	enableFailover := os.Getenv("SUMO_ENABLE_FAILOVER")
 	logTypes := os.Getenv("SUMO_LOG_TYPES")
+	enhanceJsonLogs := os.Getenv("SUMO_ENHANCE_JSON_LOGS")
 
 	if numRetry == "" {
 		cfg.NumRetry = 3
@@ -99,9 +101,11 @@ func (cfg *LambdaExtensionConfig) setDefaults() {
 		cfg.LogTypes = strings.Split(logTypes, ",")
 	}
 	if retrySleepTime == "" {
-		cfg.RetrySleepTime =  300 * time.Millisecond
+		cfg.RetrySleepTime = 300 * time.Millisecond
+	}
+	if enhanceJsonLogs == "" {
+		cfg.EnhanceJsonLogs = true
 	}
-
 }
 
 func (cfg *LambdaExtensionConfig) validateConfig() error {
@@ -111,6 +115,7 @@ func (cfg *LambdaExtensionConfig) validateConfig() error {
 	maxConcurrentRequests := os.Getenv("SUMO_MAX_CONCURRENT_REQUESTS")
 	enableFailover := os.Getenv("SUMO_ENABLE_FAILOVER")
 	retrySleepTime := os.Getenv("SUMO_RETRY_SLEEP_TIME_MS")
+	enhanceJsonLogs := os.Getenv("SUMO_ENHANCE_JSON_LOGS")
 
 	var allErrors []string
 	var err error
@@ -189,6 +194,13 @@ func (cfg *LambdaExtensionConfig) validateConfig() error {
 
 	}
 
+	if enhanceJsonLogs != "" {
+		cfg.EnhanceJsonLogs, err = strconv.ParseBool(enhanceJsonLogs)
+		if err != nil {
+			allErrors = append(allErrors, fmt.Sprintf("Unable to parse SUMO_ENHANCE_JSON_LOGS: %v", err))
+		}
+	}
+
 	// test valid log format type
 	for _, logType := range cfg.LogTypes {
 		if !utils.StringInSlice(strings.TrimSpace(logType), validLogTypes) {
diff --git a/lambda-extensions/sumoclient/sumoclient.go b/lambda-extensions/sumoclient/sumoclient.go
index 4292667..b08b6dd 100644
--- a/lambda-extensions/sumoclient/sumoclient.go
+++ b/lambda-extensions/sumoclient/sumoclient.go
@@ -202,7 +202,11 @@ func (s *sumoLogicClient) enhanceLogs(msg responseBody) {
 			if err != nil {
 				item["message"] = message
 			} else {
-				item["message"] = json
+				if s.config.EnhanceJsonLogs {
+					item["message"] = json
+				} else {
+					item = json
+				}
 			}
 		} else if ok && logType == "platform.report" {
 			s.createCWLogLine(item)
@@ -287,7 +291,7 @@ func (s *sumoLogicClient) SendLogs(ctx context.Context, rawmsg []byte) error {
 }
 
 func (s *sumoLogicClient) SendAllLogs(ctx context.Context, allMessages [][]byte) error {
-	if (len(allMessages) == 0) {
+	if len(allMessages) == 0 {
 		s.logger.Debugf("SendAllLogs: No messages to send")
 		return nil
 	}