fix(otelgorm): restore original context after "after" callback

otelgorm/internal/e2etest/e2e_test.go

@@ -93,6 +93,22 @@ func TestEndToEnd(t *testing.T) {
 				require.Equal(t, 1, len(spans))
 			},
 		},
+		{
+			do: func(ctx context.Context, db *gorm.DB) {
+				var count int64
+				query := db.WithContext(ctx).Table("generate_series(1, 10)")
+				_, _ = query.Select("*").Rows()
+				_ = query.Count(&count)
+			},
+			require: func(t *testing.T, spans []sdktrace.ReadOnlySpan) {
+				require.Equal(t, 2, len(spans))
+				require.Equal(
+					t,
+					spans[0].Parent().SpanID().String(),
+					spans[1].Parent().SpanID().String(),
+				)
+			},
+		},
 	}
 
 	for i, test := range tests {

otelgorm/otelgorm.go

@@ -1,6 +1,7 @@
 package otelgorm
 
 import (
+	"context"
 	"database/sql"
 	"database/sql/driver"
 	"fmt"
@@ -95,12 +96,18 @@ func (p otelPlugin) Initialize(db *gorm.DB) (err error) {
 	return firstErr
 }
 
+type originalCtxType int
+
+const originalCtxValue originalCtxType = iota
+
 func (p *otelPlugin) before(spanName string) gormHookFunc {
 	return func(tx *gorm.DB) {
 		if tx.DryRun && !p.includeDryRunSpans {
 			return
 		}
-		tx.Statement.Context, _ = p.tracer.Start(tx.Statement.Context, spanName, trace.WithSpanKind(trace.SpanKindClient))
+		originalCtx := tx.Statement.Context
+		newCtx := context.WithValue(originalCtx, originalCtxValue, originalCtx)
+		tx.Statement.Context, _ = p.tracer.Start(newCtx, spanName, trace.WithSpanKind(trace.SpanKindClient))
 	}
 }
 
@@ -154,6 +161,11 @@ func (p *otelPlugin) after() gormHookFunc {
 			span.RecordError(tx.Error)
 			span.SetStatus(codes.Error, tx.Error.Error())
 		}
+
+		switch originalCtx := tx.Statement.Context.Value(originalCtxValue).(type) {
+		case context.Context:
+			tx.Statement.Context = originalCtx
+		}
 	}
 }