Restore previous context in after hooks (#50)

* Try recovering previous context

* Clean up

* Run gofmt

* Add context restoration tests

Author: OrkhanAlikhanov

tracing/tracing.go

@@ -1,6 +1,7 @@
 package tracing
 
 import (
+	"context"
 	"database/sql"
 	"database/sql/driver"
 	"fmt"
@@ -108,10 +109,16 @@ func (p otelPlugin) Initialize(db *gorm.DB) (err error) {
 	return firstErr
 }
 
+type contextWrapper struct {
+	context.Context
+	parent context.Context
+}
+
 func (p *otelPlugin) before(spanName string) gormHookFunc {
 	return func(tx *gorm.DB) {
+		parentCtx := tx.Statement.Context
 		ctx, span := p.tracer.Start(tx.Statement.Context, spanName, trace.WithSpanKind(trace.SpanKindClient))
-		tx.Statement.Context = ctx
+		tx.Statement.Context = contextWrapper{ctx, parentCtx}
 
 		if !p.excludeServerAddress {
 			// `server.address` is required in the latest semconv
@@ -141,6 +148,13 @@ func (p *otelPlugin) before(spanName string) gormHookFunc {
 
 func (p *otelPlugin) after() gormHookFunc {
 	return func(tx *gorm.DB) {
+		defer func() {
+			if c, ok := tx.Statement.Context.(contextWrapper); ok {
+				// recover previous context
+				tx.Statement.Context = c.parent
+			}
+		}()
+
 		span := trace.SpanFromContext(tx.Statement.Context)
 		if !span.IsRecording() {
 			return

tracing/tracing_test.go

@@ -12,6 +12,7 @@ import (
 	"go.opentelemetry.io/otel/sdk/trace/tracetest"
 	semconv "go.opentelemetry.io/otel/semconv/v1.30.0"
 	"go.opentelemetry.io/otel/trace"
+	"go.opentelemetry.io/otel/trace/noop"
 	"gorm.io/driver/postgres"
 	"gorm.io/driver/sqlite"
 	"gorm.io/gorm"
@@ -167,6 +168,41 @@ func TestOtel(t *testing.T) {
 	}
 }
 
+func TestOtelPlugin_ContextRestoration(t *testing.T) {
+	provider := noop.NewTracerProvider()
+	p := &otelPlugin{provider: provider}
+	p.tracer = provider.Tracer("test")
+
+	origCtx := context.WithValue(context.Background(), "foo", "bar")
+	db := &gorm.DB{
+		Statement: &gorm.Statement{Context: origCtx},
+		Config:    &gorm.Config{},
+	}
+
+	// before should wrap context
+	before := p.before("test-span")
+	before(db)
+	cw, ok := db.Statement.Context.(contextWrapper)
+	require.True(t, ok)
+	require.Equal(t, origCtx, cw.parent)
+
+	// after should restore context
+	after := p.after()
+	after(db)
+	require.Equal(t, origCtx, db.Statement.Context)
+
+	origCtx = context.Background()
+	db = &gorm.DB{
+		Statement: &gorm.Statement{Context: origCtx},
+		Config:    &gorm.Config{},
+	}
+
+	// after should not panic if context is not a contextWrapper
+	after = p.after()
+	require.NotPanics(t, func() { after(db) })
+	require.Equal(t, origCtx, db.Statement.Context)
+}
+
 func attrMap(attrs []attribute.KeyValue) map[attribute.Key]attribute.Value {
 	m := make(map[attribute.Key]attribute.Value, len(attrs))
 	for _, kv := range attrs {