From 044ba475223f0818698eff9a1ba8ad2e612cc4dc Mon Sep 17 00:00:00 2001
From: urso <steffen@xata.io>
Date: Mon, 25 Oct 2021 12:24:53 +0200
Subject: [PATCH] Add zerologadapter.NewContextLogger

This change introduces a new zerologadapter that allows
users to pass the actual logger via context.Context. Especially HTTP
middleware might choose to use `(*zerolog.Logger).WithContext` and
`zerolog.Ctx`. Allowing users to extract the logger from the context
keeps the full enriched logger available when pgx emits logs.
---
 log/zerologadapter/adapter.go      | 54 +++++++++++++++++++++++-------
 log/zerologadapter/adapter_test.go | 21 ++++++++++--
 2 files changed, 61 insertions(+), 14 deletions(-)

diff --git a/log/zerologadapter/adapter.go b/log/zerologadapter/adapter.go
index 9bfa5788..6e8b4b94 100644
--- a/log/zerologadapter/adapter.go
+++ b/log/zerologadapter/adapter.go
@@ -9,9 +9,10 @@ import (
 )
 
 type Logger struct {
-	logger     zerolog.Logger
-	withFunc   func(context.Context, zerolog.Context) zerolog.Context
-	skipModule bool
+	logger      zerolog.Logger
+	withFunc    func(context.Context, zerolog.Context) zerolog.Context
+	fromContext bool
+	skipModule  bool
 }
 
 // option options for configuring the logger when creating a new logger.
@@ -38,15 +39,30 @@ 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()
-	}
+	l.init(options)
 	return &l
 }
 
+// NewContextLogger creates logger that extracts the zerolog.Logger from the
+// context.Context by using `zerolog.Ctx`. The zerolog.DefaultContextLogger will
+// be used if no logger is associated with the context.
+func NewContextLogger(options ...option) *Logger {
+	l := Logger{
+		fromContext: true,
+	}
+	l.init(options)
+	return &l
+}
+
+func (pl *Logger) init(options []option) {
+	for _, opt := range options {
+		opt(pl)
+	}
+	if !pl.skipModule {
+		pl.logger = pl.logger.With().Str("module", "pgx").Logger()
+	}
+}
+
 func (pl *Logger) Log(ctx context.Context, level pgx.LogLevel, msg string, data map[string]interface{}) {
 	var zlevel zerolog.Level
 	switch level {
@@ -63,10 +79,24 @@ func (pl *Logger) Log(ctx context.Context, level pgx.LogLevel, msg string, data
 	default:
 		zlevel = zerolog.DebugLevel
 	}
-	zctx := pl.logger.With()
+
+	var zctx zerolog.Context
+	if pl.fromContext {
+		logger := zerolog.Ctx(ctx)
+		zctx = logger.With()
+	} else {
+		zctx = pl.logger.With()
+	}
 	if pl.withFunc != nil {
 		zctx = pl.withFunc(ctx, zctx)
 	}
-	pgxlog := zctx.Fields(data).Logger()
-	pgxlog.WithLevel(zlevel).Msg(msg)
+
+	pgxlog := zctx.Logger()
+	event := pgxlog.WithLevel(zlevel)
+	if event.Enabled() {
+		if pl.fromContext && !pl.skipModule {
+			event.Str("module", "pgx")
+		}
+		event.Fields(data).Msg(msg)
+	}
 }
diff --git a/log/zerologadapter/adapter_test.go b/log/zerologadapter/adapter_test.go
index deda6ed4..3a11cbc0 100644
--- a/log/zerologadapter/adapter_test.go
+++ b/log/zerologadapter/adapter_test.go
@@ -38,6 +38,23 @@ func TestLogger(t *testing.T) {
 		}
 	})
 
+	t.Run("from context", func(t *testing.T) {
+		var buf bytes.Buffer
+		zlogger := zerolog.New(&buf)
+		ctx := zlogger.WithContext(context.Background())
+		logger := zerologadapter.NewContextLogger()
+		logger.Log(ctx, 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.Log(got)
+			t.Log(want)
+			t.Errorf("%s != %s", got, want)
+		}
+	})
+
 	var buf bytes.Buffer
 	type key string
 	var ck key
@@ -52,7 +69,8 @@ func TestLogger(t *testing.T) {
 				logWith = logWith.Str("req_id", id)
 			}
 			return logWith
-		}))
+		}),
+	)
 
 	t.Run("no request id", func(t *testing.T) {
 		buf.Reset()
@@ -77,5 +95,4 @@ func TestLogger(t *testing.T) {
 			t.Errorf("%s != %s", got, want)
 		}
 	})
-
 }