add tags to the logs

telemetry/client.go

@@ -9,6 +9,7 @@ import (
 	"log"
 	"net"
 	"net/http"
+	"os"
 	"strings"
 	"time"
 
@@ -283,13 +284,40 @@ func (c *Client) SendFunctionLogs(ctx context.Context, invokedFunctionARN string
 	return nil
 }
 
+// getNewRelicTags adds tags to the logs if NR_TAGS has values
+func getNewRelicTags(common map[string]interface{}) {
+    nrTagsStr := os.Getenv("NR_TAGS")
+    nrDelimiter := os.Getenv("NR_ENV_DELIMITER")
+    if nrDelimiter == "" {
+        nrDelimiter = ";"
+    }
+
+    if nrTagsStr != "" {
+        tags := strings.Split(nrTagsStr, nrDelimiter)
+        nrTags := make(map[string]string)
+        for _, tag := range tags {
+			fmt.Println(tag)
+            keyValue := strings.Split(tag, ":")
+            if len(keyValue) == 2 {
+                nrTags[keyValue[0]] = keyValue[1]
+            }
+        }
+
+        for k, v := range nrTags {
+            common[k] = v
+        }
+    }
+}
+
 // buildLogPayloads is a helper function that improves readability of the SendFunctionLogs method
 func (c *Client) buildLogPayloads(ctx context.Context, invokedFunctionARN string, lines []logserver.LogLine) ([]*bytes.Buffer, requestBuilder, error) {
 	common := map[string]interface{}{
 		"plugin":    util.Id,
 		"faas.arn":  invokedFunctionARN,
 		"faas.name": c.functionName,
 	}
+
+	getNewRelicTags(common)
 
 	logMessages := make([]FunctionLogMessage, 0, len(lines))
 	for _, l := range lines {

telemetry/client_test.go

@@ -6,6 +6,7 @@ import (
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"os"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -462,3 +463,104 @@ func TestGetLogEndpointURL(t *testing.T) {
 	assert.Equal(t, LogEndpointUS, getLogEndpointURL("us mock license key", ""))
 	assert.Equal(t, LogEndpointEU, getLogEndpointURL("eu mock license key", ""))
 }
+func TestGetNewRelicTags(t *testing.T) {
+    os.Setenv("NR_TAGS", "env:prod;team:myTeam")
+    os.Setenv("NR_ENV_DELIMITER", ";")
+
+    tests := []struct {
+        name        string
+        common      map[string]interface{}
+        expected    map[string]interface{}
+        envTags     string
+        envDelimiter string
+    }{
+        {
+            name: "Add New Relic tags to common",
+            common: map[string]interface{}{
+                "plugin":    "testPlugin",
+                "faas.arn":  "arn:aws:lambda:us-east-1:123456789012:function:testFunction",
+                "faas.name": "testFunction",
+            },
+            expected: map[string]interface{}{
+                "plugin":    "testPlugin",
+                "faas.arn":  "arn:aws:lambda:us-east-1:123456789012:function:testFunction",
+                "faas.name": "testFunction",
+                "env":       "prod",
+                "team":      "myTeam",
+            },
+            envTags:     "env:prod;team:myTeam",
+            envDelimiter: ";",
+        },
+		{
+            name: "Add New Relic tags to common if no delimiter is set",
+            common: map[string]interface{}{
+                "plugin":    "testPlugin",
+                "faas.arn":  "arn:aws:lambda:us-east-1:123456789012:function:testFunction",
+                "faas.name": "testFunction",
+            },
+            expected: map[string]interface{}{
+                "plugin":    "testPlugin",
+                "faas.arn":  "arn:aws:lambda:us-east-1:123456789012:function:testFunction",
+                "faas.name": "testFunction",
+                "env":       "prod",
+                "team":      "myTeam",
+            },
+            envTags:     "env:prod;team:myTeam",
+        },
+		{
+            name: "No New Relic tags to common if delimiter is incorrect",
+            common: map[string]interface{}{
+                "plugin":    "testPlugin",
+                "faas.arn":  "arn:aws:lambda:us-east-1:123456789012:function:testFunction",
+                "faas.name": "testFunction",
+            },
+            expected: map[string]interface{}{
+                "plugin":    "testPlugin",
+                "faas.arn":  "arn:aws:lambda:us-east-1:123456789012:function:testFunction",
+                "faas.name": "testFunction",
+            },
+            envTags:     "env:prod;team:myTeam",
+            envDelimiter: ",",
+        },
+        {
+            name: "No New Relic tags to add",
+            common: map[string]interface{}{
+                "plugin":    "testPlugin",
+                "faas.arn":  "arn:aws:lambda:us-east-1:123456789012:function:testFunction",
+                "faas.name": "testFunction",
+            },
+            expected: map[string]interface{}{
+                "plugin":    "testPlugin",
+                "faas.arn":  "arn:aws:lambda:us-east-1:123456789012:function:testFunction",
+                "faas.name": "testFunction",
+            },
+            envTags:     "",
+            envDelimiter: "",
+        },
+    }
+
+    for _, tt := range tests {
+        t.Run(tt.name, func(t *testing.T) {
+            os.Setenv("NR_TAGS", tt.envTags)
+            os.Setenv("NR_ENV_DELIMITER", tt.envDelimiter)
+
+            common := make(map[string]interface{}, len(tt.common))
+            for k, v := range tt.common {
+                common[k] = v
+            }
+
+            getNewRelicTags(common)
+
+            for k, v := range tt.expected {
+                if common[k] != v {
+                    t.Errorf("expected common[%q] to be %v, but got %v", k, v, common[k])
+                }
+            }
+            for k := range common {
+                if _, ok := tt.expected[k]; !ok {
+                    t.Errorf("unexpected key %q in common map", k)
+                }
+            }
+        })
+    }
+}