From 509684e614b037ee490f04c9fc8c5d08b9e57d61 Mon Sep 17 00:00:00 2001
From: lanzafame <adrianlanzafame92@gmail.com>
Date: Thu, 25 Feb 2021 13:14:10 +1000
Subject: [PATCH] feat: add ability to specify labels for all loggers

---
 setup.go      | 21 +++++++++++++++++++++
 setup_test.go | 34 +++++++++++++++++++++++++++++++++-
 2 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/setup.go b/setup.go
index 9820cf0..1f97e8c 100644
--- a/setup.go
+++ b/setup.go
@@ -33,6 +33,7 @@ const (
 	envLoggingURL = "GOLOG_URL" // url that will be processed by sink in the zap
 
 	envLoggingOutput = "GOLOG_OUTPUT" // possible values: stdout|stderr|file combine multiple values with '+'
+	envLoggingLabels = "GOLOG_LOG_LABELS" // comma-separated key-value pairs, i.e. "app,example_app,dc,sjc-1"
 )
 
 type LogFormat int
@@ -61,6 +62,9 @@ type Config struct {
 
 	// URL with schema supported by zap. Use zap.RegisterSink
 	URL string
+
+	// Labels is a set of key-values to apply to all loggers
+	Labels []string
 }
 
 // ErrNoSuchLogger is returned when the util pkg is asked for a non existant logger
@@ -122,6 +126,13 @@ func SetupLogging(cfg Config) {
 	}
 
 	newPrimaryCore := newCore(primaryFormat, ws, LevelDebug) // the main core needs to log everything.
+
+	if len(cfg.Labels) > 0 {
+		for i := 0; i < len(cfg.Labels); i=i+2 {
+			newPrimaryCore = newPrimaryCore.With([]zap.Field{zap.String(cfg.Labels[i], cfg.Labels[i+1])})
+		}
+	}
+
 	if primaryCore != nil {
 		loggerCore.ReplaceCore(primaryCore, newPrimaryCore)
 	} else {
@@ -293,5 +304,15 @@ func configFromEnv() Config {
 		}
 	}
 
+	labels := os.Getenv(envLoggingLabels)
+	if labels != "" {
+		labelKVs := strings.Split(labels, ",")
+		if len(labelKVs)%2 != 0 {
+			fmt.Fprint(os.Stderr, "odd number of args for GOLOG_LOG_LABELS; please specify complete key-value pairs")
+		} else {
+			cfg.Labels = labelKVs
+		}
+	}
+
 	return cfg
 }
diff --git a/setup_test.go b/setup_test.go
index 549e089..313f956 100644
--- a/setup_test.go
+++ b/setup_test.go
@@ -121,4 +121,36 @@ func TestLogToFile(t *testing.T) {
 		t.Logf("want: '%s', got: '%s'", want, string(content))
 		t.Fail()
 	}
-}
\ No newline at end of file
+}
+
+func TestLogLabels(t *testing.T) {
+	r, w, err := os.Pipe()
+	if err != nil {
+		t.Fatalf("failed to open pipe: %v", err)
+	}
+
+	stderr := os.Stderr
+	os.Stderr = w
+	defer func() {
+		os.Stderr = stderr
+	}()
+
+	// set the go-log labels env var
+	os.Setenv(envLoggingLabels, "app,example_app,dc,sjc-1")
+	SetupLogging(configFromEnv())
+
+	log := getLogger("test")
+
+	log.Error("scooby")
+	w.Close()
+
+	buf := &bytes.Buffer{}
+	if _, err := io.Copy(buf, r); err != nil && err != io.ErrClosedPipe {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	t.Log(buf.String())
+	if !strings.Contains(buf.String(), "{\"app\": \"example_app\", \"dc\": \"sjc-1\"}") {
+		t.Errorf("got %q, wanted it to contain log output", buf.String())
+	}
+}