From a594aefccdb756b98672371ab7d242463dbe1221 Mon Sep 17 00:00:00 2001
From: Khushi Jain <khushij393@gmail.com>
Date: Tue, 4 Jun 2024 13:05:01 +0530
Subject: [PATCH] otelzap: Allow context injection via fields (#5707)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Part of
https://github.com/open-telemetry/opentelemetry-go-contrib/issues/5191

Pre-work
https://github.com/open-telemetry/opentelemetry-go-contrib/pull/5279

---------

Co-authored-by: Robert Pająk <pellared@hotmail.com>
---
 bridges/otelzap/core.go      | 28 +++++++++++++++++++++-------
 bridges/otelzap/core_test.go |  3 +++
 2 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/bridges/otelzap/core.go b/bridges/otelzap/core.go
index e4be7f783b8..f90be80de70 100644
--- a/bridges/otelzap/core.go
+++ b/bridges/otelzap/core.go
@@ -90,6 +90,7 @@ func WithLoggerProvider(provider log.LoggerProvider) Option {
 type Core struct {
 	logger log.Logger
 	attr   []log.KeyValue
+	ctx    context.Context
 }
 
 // Compile-time check *Core implements zapcore.Core.
@@ -100,6 +101,7 @@ func NewCore(name string, opts ...Option) *Core {
 	cfg := newConfig(opts)
 	return &Core{
 		logger: cfg.logger(name),
+		ctx:    context.Background(),
 	}
 }
 
@@ -114,7 +116,11 @@ func (o *Core) Enabled(level zapcore.Level) bool {
 func (o *Core) With(fields []zapcore.Field) zapcore.Core {
 	cloned := o.clone()
 	if len(fields) > 0 {
-		cloned.attr = append(cloned.attr, convertField(fields)...)
+		ctx, attrbuf := convertField(fields)
+		if ctx != nil {
+			cloned.ctx = ctx
+		}
+		cloned.attr = append(cloned.attr, attrbuf...)
 	}
 	return cloned
 }
@@ -123,6 +129,7 @@ func (o *Core) clone() *Core {
 	return &Core{
 		logger: o.logger,
 		attr:   slices.Clone(o.attr),
+		ctx:    o.ctx,
 	}
 }
 
@@ -148,28 +155,35 @@ func (o *Core) Write(ent zapcore.Entry, fields []zapcore.Field) error {
 	r.SetBody(log.StringValue(ent.Message))
 	r.SetSeverity(convertLevel(ent.Level))
 
-	// TODO: Handle attributes passed via With (exceptions: context.Context and zap.Namespace).
-	// TODO: Handle context.Context containing trace context.
 	// TODO: Handle zap.Namespace.
 	// TODO: Handle ent.LoggerName.
 
 	r.AddAttributes(o.attr...)
 	if len(fields) > 0 {
-		r.AddAttributes(convertField(fields)...)
+		ctx, attrbuf := convertField(fields)
+		if ctx != nil {
+			o.ctx = ctx
+		}
+		r.AddAttributes(attrbuf...)
 	}
 
-	o.logger.Emit(context.Background(), r)
+	o.logger.Emit(o.ctx, r)
 	return nil
 }
 
-func convertField(fields []zapcore.Field) []log.KeyValue {
+func convertField(fields []zapcore.Field) (context.Context, []log.KeyValue) {
 	// TODO: Use objectEncoder from a pool instead of newObjectEncoder.
+	var ctx context.Context
 	enc := newObjectEncoder(len(fields))
 	for _, field := range fields {
+		if ctxFld, ok := field.Interface.(context.Context); ok {
+			ctx = ctxFld
+			continue
+		}
 		field.AddTo(enc)
 	}
 
-	return enc.kv
+	return ctx, enc.kv
 }
 
 func convertLevel(level zapcore.Level) log.Severity {
diff --git a/bridges/otelzap/core_test.go b/bridges/otelzap/core_test.go
index 82e6a23b8a6..8dea6323826 100644
--- a/bridges/otelzap/core_test.go
+++ b/bridges/otelzap/core_test.go
@@ -44,6 +44,9 @@ func TestCore(t *testing.T) {
 
 	rec.Reset()
 
+	// TODO: Add WriteContext test case.
+	// TODO: Add WithContext test case.
+
 	// test child logger with accumulated fields
 	t.Run("With", func(t *testing.T) {
 		testCases := [][]string{{"test1", "value1"}, {"test2", "value2"}}