add tests for trace context propagation

README.md

@@ -297,7 +297,7 @@ func main() {
 }
 
 ```
-Trace context propogation using [OTEL SDK v1.7.0](https://github.com/open-telemetry/opentelemetry-go/tree/v1.7.0)
+Trace context propogation using OTEL SDK
 ```go
 package main
 

conn_test.go

@@ -13,6 +13,10 @@ import (
 
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/suite"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/propagation"
+	"go.opentelemetry.io/otel/sdk/trace"
+	oteltrace "go.opentelemetry.io/otel/trace"
 )
 
 var (
@@ -377,6 +381,72 @@ func (s *connSuite) TestBuildRequestWithQueryIdAndQuotaKey() {
 		}
 	}
 }
+
+func (s *connSuite) TestBuildRequestWithTraceContext() {
+	cn := newConn(NewConfig())
+	testCases := []struct {
+		name                       string
+		traceProvider              oteltrace.TracerProvider
+		propagator                 propagation.TextMapPropagator
+		expectedValidSpan          bool
+		expectedTraceHeaderPresent bool
+	}{
+		{
+			name:                       "trace context",
+			traceProvider:              trace.NewTracerProvider(),
+			propagator:                 propagation.TraceContext{},
+			expectedTraceHeaderPresent: true,
+			expectedValidSpan:          true,
+		},
+		{
+			name:                       "trace context with noop trace provider",
+			traceProvider:              oteltrace.NewNoopTracerProvider(),
+			propagator:                 propagation.TraceContext{},
+			expectedTraceHeaderPresent: false,
+			expectedValidSpan:          false,
+		},
+		{
+			name:                       "trace context with noop propagator",
+			traceProvider:              trace.NewTracerProvider(),
+			propagator:                 propagation.NewCompositeTextMapPropagator(),
+			expectedTraceHeaderPresent: false,
+			expectedValidSpan:          false,
+		},
+		{
+			name:                       "trace context with noop provider and propagator",
+			traceProvider:              oteltrace.NewNoopTracerProvider(),
+			propagator:                 propagation.NewCompositeTextMapPropagator(),
+			expectedTraceHeaderPresent: false,
+			expectedValidSpan:          false,
+		},
+	}
+	for _, tc := range testCases {
+		ctx := context.Background()
+		otel.SetTextMapPropagator(tc.propagator)
+		// start span
+		trCtx, expectedSpan := tc.traceProvider.Tracer("go test").Start(ctx, tc.name)
+		req, err := cn.buildRequest(trCtx, "SELECT 1", nil)
+		expectedSpan.End()
+		if s.NoError(err) {
+			// check TraceParent header present in request header
+			traceHeader := req.Header.Get("TraceParent")
+			traceHeaderPresent := len(traceHeader) > 0
+			s.Equal(tc.expectedTraceHeaderPresent, traceHeaderPresent)
+
+			reqCtx := context.Background()
+			reqTrCtx := otel.GetTextMapPropagator().Extract(reqCtx, propagation.HeaderCarrier(req.Header))
+			reqSpan := oteltrace.SpanFromContext(reqTrCtx).SpanContext()
+
+			// check if request span valid
+			s.Equal(tc.expectedValidSpan, reqSpan.IsValid())
+			if tc.expectedValidSpan {
+				s.Equal(expectedSpan.SpanContext().SpanID(), reqSpan.SpanID())
+				s.Equal(expectedSpan.SpanContext().TraceID(), reqSpan.TraceID())
+			}
+		}
+	}
+}
+
 func (s *connSuite) TestBuildRequestParamsInterpolation() {
 	query := `INSERT INTO test (str) VALUES ("Question?")`
 	cn := newConn(NewConfig())

go.mod

@@ -4,6 +4,7 @@ require (
 	github.com/google/uuid v1.2.0
 	github.com/stretchr/testify v1.7.1
 	go.opentelemetry.io/otel v1.7.0
+	go.opentelemetry.io/otel/sdk v1.7.0
 	go.opentelemetry.io/otel/trace v1.7.0
 )
 

go.sum

@@ -16,8 +16,12 @@ github.com/stretchr/testify v1.7.1 h1:5TQK59W5E3v0r2duFAb7P95B6hEeOyEnHRa8MjYSMT
 github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 go.opentelemetry.io/otel v1.7.0 h1:Z2lA3Tdch0iDcrhJXDIlC94XE+bxok1F9B+4Lz/lGsM=
 go.opentelemetry.io/otel v1.7.0/go.mod h1:5BdUoMIz5WEs0vt0CUEMtSSaTSHBBVwrhnz7+nrD5xk=
+go.opentelemetry.io/otel/sdk v1.7.0 h1:4OmStpcKVOfvDOgCt7UriAPtKolwIhxpnSNI/yK+1B0=
+go.opentelemetry.io/otel/sdk v1.7.0/go.mod h1:uTEOTwaqIVuTGiJN7ii13Ibp75wJmYUDe374q6cZwUU=
 go.opentelemetry.io/otel/trace v1.7.0 h1:O37Iogk1lEkMRXewVtZ1BBTVn5JEp8GrJvP92bJqC6o=
 go.opentelemetry.io/otel/trace v1.7.0/go.mod h1:fzLSB9nqR2eXzxPXb2JW9IKE+ScyXA48yyE4TNvoHqU=
+golang.org/x/sys v0.0.0-20210423185535-09eb48e85fd7 h1:iGu644GcxtEcrInvDsQRCwJjtCIOlT2V7IRt6ah2Whw=
+golang.org/x/sys v0.0.0-20210423185535-09eb48e85fd7/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=