set level

ITV · Jun 13, 2024 · 24a505e · 24a505e
1 parent c3530c4
commit 24a505e
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 12 deletions.
diff --git a/loki_extension/Taskfile.yml b/loki_extension/Taskfile.yml
@@ -19,7 +19,7 @@ tasks:
           echo $ARN && \
           LAYER_VERSION=`echo $ARN | cut -d: -f8` && \
           echo $LAYER_VERSION && \
-          aws lambda add-layer-version-permission --layer-name "grafana-loki-extension" --version-number $LAYER_VERSION --statement-id "test" --action lambda:GetLayerVersion --principal "*" --organization-id "o-38aomgzvwn" --output text
+          aws lambda add-layer-version-permission --layer-name "grafana-loki-extension" --version-number $LAYER_VERSION --statement-id "layer-perm" --action lambda:GetLayerVersion --principal "*" --organization-id "o-38aomgzvwn" --output text
     silent: false
 
   build:

diff --git a/loki_extension/go.mod b/loki_extension/go.mod
@@ -3,7 +3,7 @@ module loki-extension
 go 1.18
 
 require (
-	github.com/ITV/promtail-client-1 v0.0.0-20240612134402-26913bd53ff4
+	github.com/ITV/promtail-client-1 v0.0.0-20240613121915-db48bd16489d
 	github.com/golang-collections/go-datastructures v0.0.0-20150211160725-59788d5eb259
 	github.com/pkg/errors v0.9.1
 	github.com/sirupsen/logrus v1.9.0

diff --git a/loki_extension/go.sum b/loki_extension/go.sum
@@ -1,7 +1,7 @@
-github.com/ITV/promtail-client-1 v0.0.0-20240612134402-26913bd53ff4 h1:yKyUDd3FXgdCv96y95urb5O8UVx8oFUat+/H9KAmGrI=
-github.com/ITV/promtail-client-1 v0.0.0-20240612134402-26913bd53ff4/go.mod h1:sHaIS1I8USe2evPTw8yNea3oUSklV76sjnBiQX1CzgY=
-github.com/afiskon/promtail-client v0.0.0-20190305142237-506f3f921e9c h1:AMDVOKGaiqse4qiRXSzRgpC9DCNTHCx6zpzdtXXrKM4=
-github.com/afiskon/promtail-client v0.0.0-20190305142237-506f3f921e9c/go.mod h1:p/7Wos+jcfrnwLqqzJMZ0s323kfVtJPW+HUvAANklVQ=
+github.com/ITV/promtail-client-1 v0.0.0-20240613114054-ade2bc567bec h1:HKPLLK0CV4NlPWZmDe1/fPFS8y17v/Wq5TLBjRj4bBw=
+github.com/ITV/promtail-client-1 v0.0.0-20240613114054-ade2bc567bec/go.mod h1:sHaIS1I8USe2evPTw8yNea3oUSklV76sjnBiQX1CzgY=
+github.com/ITV/promtail-client-1 v0.0.0-20240613121915-db48bd16489d h1:+uTgcYZuuDt22V9yvf2fjut9wQnKrSwV6rmUojI2LvE=
+github.com/ITV/promtail-client-1 v0.0.0-20240613121915-db48bd16489d/go.mod h1:sHaIS1I8USe2evPTw8yNea3oUSklV76sjnBiQX1CzgY=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=

diff --git a/loki_extension/loki/promtail.go b/loki_extension/loki/promtail.go
@@ -1,6 +1,7 @@
 package loki
 
 import (
+	"encoding/json"
 	"fmt"
 	"log"
 	"os"
@@ -14,6 +15,7 @@ type LogEntry struct {
 	Record string `json:"record"`
 	Time   string `json:"time"`
 	Type   string `json:"type"`
+	Level  string `json:"level"`
 }
 
 var (
@@ -31,8 +33,6 @@ func init() {
 	sendLabels = []string{"source = lambda"}
 
 	// CP Standard Labels
-	// product := os.Getenv("PRODUCT")     - minimise labels: arguably redundant
-	// ecosystem := os.Getenv("ECOSYSTEM") - minimise labels: arguably redundant
 	environment := os.Getenv("ENVIRONMENT")
 	service := os.Getenv("SERVICE")
 	component := os.Getenv("COMPONENT")
@@ -65,8 +65,6 @@ func init() {
 		Labels:             fmt.Sprintf("{%s}", strings.Join(labels, ",")),
 		BatchWait:          5 * time.Second,
 		BatchEntriesNumber: 10000,
-		SendLevel:          promtail.INFO,
-		PrintLevel:         promtail.ERROR,
 	}
 	loki, err = promtail.NewClientProto(conf)
 	if err != nil {
@@ -76,8 +74,31 @@ func init() {
 }
 
 func LokiSend(record *string) {
-	tstamp := time.Now().String()
-	loki.Infof("{\"@timestamp\": \"%s\", \"message\": \"%v\"}", tstamp, strings.TrimSuffix(*record, "\n"))
+	var logEntry LogEntry
+	level := "INFO"
+	if json.Unmarshal([]byte(*record), &logEntry) == nil {
+		if logEntry.Level != "" {
+			level = logEntry.Level
+		}
+		newLabels := append(labels, fmt.Sprintf("level=\"%s\"", level))
+		conf.Labels = fmt.Sprintf("{%s}", strings.Join(newLabels, ","))
+		loki, err = promtail.NewClientProto(conf)
+		if err != nil {
+			log.Println("Promtail re-init error")
+			log.Println(err)
+			return
+		}
+	} else {
+		newLabels := append(labels, fmt.Sprintf("level=\"%s\"", level))
+		conf.Labels = fmt.Sprintf("{%s}", strings.Join(newLabels, ","))
+		loki, err = promtail.NewClientProto(conf)
+		if err != nil {
+			log.Println("Promtail re-init error")
+			log.Println(err)
+			return
+		}
+	}
+	loki.Infof(strings.TrimSuffix(*record, "\n"))
 }
 
 func LokiShutdown() {