From 24cc151938cfd4a7477cda5919188e2397e9bbac Mon Sep 17 00:00:00 2001
From: ycc <yves.cc@laposte.net>
Date: Tue, 4 Jun 2024 22:34:44 +0200
Subject: [PATCH] lokiwriter fix

---
 lokiwriter.go | 74 ++++++++++++++++++++++++++++++++++++++-------------
 1 file changed, 56 insertions(+), 18 deletions(-)

diff --git a/lokiwriter.go b/lokiwriter.go
index b680119..d6de938 100644
--- a/lokiwriter.go
+++ b/lokiwriter.go
@@ -5,8 +5,10 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
-	"strings"
+	"os"
 	"time"
+
+	"github.com/rs/zerolog"
 )
 
 type LokiWriter struct {
@@ -20,13 +22,8 @@ type LokiPayload struct {
 }
 
 type LokiStream struct {
-	Labels  string      `json:"labels"`
-	Entries []LokiEntry `json:"entries"`
-}
-
-type LokiEntry struct {
-	Timestamp string `json:"ts"`
-	Line      string `json:"line"`
+	Stream map[string]string `json:"stream"`
+	Values [][]string        `json:"values"`
 }
 
 func NewLokiWriter(url string, labels map[string]string) *LokiWriter {
@@ -38,21 +35,37 @@ func NewLokiWriter(url string, labels map[string]string) *LokiWriter {
 }
 
 func (w *LokiWriter) Write(p []byte) (n int, err error) {
-	// Format log entry for Loki
-	entry := LokiEntry{
-		Timestamp: time.Now().Format(time.RFC3339Nano),
-		Line:      string(p),
+	// Use zerolog to parse the log level
+	var event map[string]interface{}
+	if err := json.Unmarshal(p, &event); err != nil {
+		return 0, fmt.Errorf("failed to unmarshal log event: %w", err)
 	}
 
-	labels := []string{}
-	for k, v := range w.labels {
-		labels = append(labels, fmt.Sprintf(`%s="%s"`, k, v))
+	level := ""
+	if l, ok := event["level"].(string); ok {
+		level = l
 	}
-	labelString := "{" + strings.Join(labels, ",") + "}"
+
+	message := ""
+	if m, ok := event["message"].(string); ok {
+		message = m
+	}
+
+	// Add log level to labels
+	labels := make(map[string]string)
+	for k, v := range w.labels {
+		labels[k] = v
+	}
+	labels["level"] = level
+
+	// Format the timestamp in nanoseconds
+	timestamp := fmt.Sprintf("%d000000", time.Now().UnixNano()/int64(time.Millisecond))
 
 	stream := LokiStream{
-		Labels:  labelString,
-		Entries: []LokiEntry{entry},
+		Stream: labels,
+		Values: [][]string{
+			{timestamp, message},
+		},
 	}
 
 	payload := LokiPayload{
@@ -64,6 +77,8 @@ func (w *LokiWriter) Write(p []byte) (n int, err error) {
 		return 0, fmt.Errorf("failed to marshal payload: %w", err)
 	}
 
+	fmt.Printf("Sending payload to Loki: %s\n", string(payloadBytes))
+
 	req, err := http.NewRequest("POST", w.url, bytes.NewReader(payloadBytes))
 	if err != nil {
 		return 0, fmt.Errorf("failed to create HTTP request: %w", err)
@@ -76,9 +91,32 @@ func (w *LokiWriter) Write(p []byte) (n int, err error) {
 	}
 	defer resp.Body.Close()
 
+	fmt.Printf("Loki response status: %d\n", resp.StatusCode)
 	if resp.StatusCode != http.StatusNoContent {
 		return 0, fmt.Errorf("received non-204 response from Loki: %d", resp.StatusCode)
 	}
 
 	return len(p), nil
 }
+
+func main() {
+	lokiURL := "http://your-loki-url/loki/api/v1/push"
+	labels := map[string]string{
+		"app": "your_app_name",
+		// Add more labels as needed
+	}
+
+	lokiWriter := NewLokiWriter(lokiURL, labels)
+
+	consoleWriter := zerolog.ConsoleWriter{Out: os.Stdout, TimeFormat: time.RFC3339}
+
+	multiWriter := zerolog.MultiLevelWriter(consoleWriter, lokiWriter)
+
+	logger := zerolog.New(multiWriter).With().Timestamp().Logger()
+
+	// Configure the logger in the sublibrary
+	sublibrary.AddLogger(logger)
+
+	logger.Info().Msg("Main library log message")
+	sublibrary.SublibraryFunction()
+}