From 255276c39005c663ce7b68978d2ed6a1a2e27bd0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Thomas=20Fr=C3=B6ssman?= <thomasf@jossystem.se>
Date: Fri, 17 Sep 2021 09:06:07 +0200
Subject: [PATCH] Add context options to zerologadapter

WithContextFunc adds possibility to get request scoped values from the
ctx.Context before logging lines.

WithoutPGXModule disables adding module:pgx to the default logger context.
---
 log/zerologadapter/adapter.go      | 44 +++++++++++++---
 log/zerologadapter/adapter_test.go | 81 ++++++++++++++++++++++++++++++
 2 files changed, 118 insertions(+), 7 deletions(-)
 create mode 100644 log/zerologadapter/adapter_test.go

diff --git a/log/zerologadapter/adapter.go b/log/zerologadapter/adapter.go
index efbcb5bf..9bfa5788 100644
--- a/log/zerologadapter/adapter.go
+++ b/log/zerologadapter/adapter.go
@@ -9,15 +9,42 @@ import (
 )
 
 type Logger struct {
-	logger zerolog.Logger
+	logger     zerolog.Logger
+	withFunc   func(context.Context, zerolog.Context) zerolog.Context
+	skipModule bool
+}
+
+// option options for configuring the logger when creating a new logger.
+type option func(logger *Logger)
+
+// WithContextFunc adds possibility to get request scoped values from the
+// ctx.Context before logging lines.
+func WithContextFunc(withFunc func(context.Context, zerolog.Context) zerolog.Context) option {
+	return func(logger *Logger) {
+		logger.withFunc = withFunc
+	}
+}
+
+// WithoutPGXModule disables adding module:pgx to the default logger context.
+func WithoutPGXModule() option {
+	return func(logger *Logger) {
+		logger.skipModule = true
+	}
 }
 
 // NewLogger accepts a zerolog.Logger as input and returns a new custom pgx
-// logging fascade as output.
-func NewLogger(logger zerolog.Logger) *Logger {
-	return &Logger{
-		logger: logger.With().Str("module", "pgx").Logger(),
+// logging facade as output.
+func NewLogger(logger zerolog.Logger, options ...option) *Logger {
+	l := Logger{
+		logger: logger,
 	}
+	for _, opt := range options {
+		opt(&l)
+	}
+	if !l.skipModule {
+		l.logger = l.logger.With().Str("module", "pgx").Logger()
+	}
+	return &l
 }
 
 func (pl *Logger) Log(ctx context.Context, level pgx.LogLevel, msg string, data map[string]interface{}) {
@@ -36,7 +63,10 @@ func (pl *Logger) Log(ctx context.Context, level pgx.LogLevel, msg string, data
 	default:
 		zlevel = zerolog.DebugLevel
 	}
-
-	pgxlog := pl.logger.With().Fields(data).Logger()
+	zctx := pl.logger.With()
+	if pl.withFunc != nil {
+		zctx = pl.withFunc(ctx, zctx)
+	}
+	pgxlog := zctx.Fields(data).Logger()
 	pgxlog.WithLevel(zlevel).Msg(msg)
 }
diff --git a/log/zerologadapter/adapter_test.go b/log/zerologadapter/adapter_test.go
new file mode 100644
index 00000000..deda6ed4
--- /dev/null
+++ b/log/zerologadapter/adapter_test.go
@@ -0,0 +1,81 @@
+package zerologadapter_test
+
+import (
+	"bytes"
+	"context"
+	"testing"
+
+	"github.com/jackc/pgx/v4"
+	"github.com/jackc/pgx/v4/log/zerologadapter"
+	"github.com/rs/zerolog"
+)
+
+func TestLogger(t *testing.T) {
+
+	t.Run("default", func(t *testing.T) {
+		var buf bytes.Buffer
+		zlogger := zerolog.New(&buf)
+		logger := zerologadapter.NewLogger(zlogger)
+		logger.Log(context.Background(), pgx.LogLevelInfo, "hello", map[string]interface{}{"one": "two"})
+		const want = `{"level":"info","module":"pgx","one":"two","message":"hello"}
+`
+		got := buf.String()
+		if got != want {
+			t.Errorf("%s != %s", got, want)
+		}
+	})
+
+	t.Run("disable pgx module", func(t *testing.T) {
+		var buf bytes.Buffer
+		zlogger := zerolog.New(&buf)
+		logger := zerologadapter.NewLogger(zlogger, zerologadapter.WithoutPGXModule())
+		logger.Log(context.Background(), pgx.LogLevelInfo, "hello", nil)
+		const want = `{"level":"info","message":"hello"}
+`
+		got := buf.String()
+		if got != want {
+			t.Errorf("%s != %s", got, want)
+		}
+	})
+
+	var buf bytes.Buffer
+	type key string
+	var ck key
+	zlogger := zerolog.New(&buf)
+	logger := zerologadapter.NewLogger(zlogger,
+		zerologadapter.WithContextFunc(func(ctx context.Context, logWith zerolog.Context) zerolog.Context {
+			// You can use zerolog.hlog.IDFromCtx(ctx) or even
+			// zerolog.log.Ctx(ctx) to fetch the whole logger instance from the
+			// context if you want.
+			id, ok := ctx.Value(ck).(string)
+			if ok {
+				logWith = logWith.Str("req_id", id)
+			}
+			return logWith
+		}))
+
+	t.Run("no request id", func(t *testing.T) {
+		buf.Reset()
+		ctx := context.Background()
+		logger.Log(ctx, pgx.LogLevelInfo, "hello", nil)
+		const want = `{"level":"info","module":"pgx","message":"hello"}
+`
+		got := buf.String()
+		if got != want {
+			t.Errorf("%s != %s", got, want)
+		}
+	})
+
+	t.Run("with request id", func(t *testing.T) {
+		buf.Reset()
+		ctx := context.WithValue(context.Background(), ck, "1")
+		logger.Log(ctx, pgx.LogLevelInfo, "hello", map[string]interface{}{"two": "2"})
+		const want = `{"level":"info","module":"pgx","req_id":"1","two":"2","message":"hello"}
+`
+		got := buf.String()
+		if got != want {
+			t.Errorf("%s != %s", got, want)
+		}
+	})
+
+}