Merge pull request #188 from pranavmehta94/pranav/tracing

DoubleDi · web-flow · commit a1b09eccc31d · 2024-11-23T02:48:17.000+09:00
Add support to send trace context to clickhouse server
diff --git a/README.md b/README.md
@@ -297,7 +297,69 @@ func main() {
 }
 
 ```
+Trace context propogation using OTEL SDK
+```go
+package main
+
+import (
+	"context"
+	"fmt"
+	"log"
+
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/propagation"
+	"go.opentelemetry.io/otel/sdk/trace"
+	oteltrace "go.opentelemetry.io/otel/trace"
+
+	"database/sql"
+
+	_ "github.com/mailru/go-clickhouse/v2"
+)
+
+func startTracing() (oteltrace.TracerProvider, error) {
+	return trace.NewTracerProvider(), nil
+}
+
+func main() {
+	// Open DB connection
+	connect, err := sql.Open("chhttp", "http://127.0.0.1:8123/default")
+	if err != nil {
+		log.Fatal(err)
+	}
+	ctx := context.Background()
+
+	// Get trace provider
+	tp, err := startTracing()
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	// Set MapPropagator
+	otel.SetTextMapPropagator(propagation.TraceContext{})
+
+	if err := connect.PingContext(ctx); err != nil {
+		log.Fatal(err)
+	}
+
+	// start span
+	trCtx, span := tp.Tracer("test").Start(ctx, "app-query")
 
+	// execute query with span context
+	rows, err := connect.QueryContext(trCtx, "SELECT COUNT() FROM (SELECT number FROM system.numbers LIMIT 5)")
+	if err != nil {
+		log.Fatal(err)
+	}
+	span.End()
+	var count uint64
+	for rows.Next() {
+		if err := rows.Scan(&count); err != nil {
+			log.Fatal(err)
+		}
+	}
+	fmt.Printf("count: %d\n", count)
+}
+
+```
 ## Go versions
 Officially support last 4 golang releases
 
diff --git a/conn.go b/conn.go
@@ -19,6 +19,9 @@ import (
 	"time"
 
 	"github.com/google/uuid"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/propagation"
+	"go.opentelemetry.io/otel/trace"
 )
 
 type key int
@@ -374,6 +377,12 @@ func (c *conn) buildRequest(ctx context.Context, query string, params []driver.V
 		req.Header.Set("Content-Encoding", "gzip")
 	}
 
+	parentSpan := trace.SpanFromContext(ctx)
+	if parentSpan.SpanContext().IsValid() {
+		carrier := propagation.HeaderCarrier(req.Header)
+		otel.GetTextMapPropagator().Inject(ctx, carrier)
+	}
+
 	return req, nil
 }
 
diff --git a/conn_test.go b/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())
diff --git a/go.mod b/go.mod
@@ -2,7 +2,10 @@ module github.com/mailru/go-clickhouse/v2
 
 require (
 	github.com/google/uuid v1.2.0
-	github.com/stretchr/testify v1.3.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 1.11
+go 1.16
diff --git a/go.sum b/go.sum
@@ -1,9 +1,29 @@
 github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A=
+github.com/go-logr/logr v1.2.3 h1:2DntVwHkVopvECVRSlL5PSo9eG+cAkDCuckLubN+rq0=
+github.com/go-logr/logr v1.2.3/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A=
+github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
+github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
+github.com/google/go-cmp v0.5.7 h1:81/ik6ipDQS2aGcBfIN5dHDB36BwrStyeAQquSYCV4o=
+github.com/google/go-cmp v0.5.7/go.mod h1:n+brtR0CgQNWTVd5ZUFpTBC8YFBDLK/h/bpaJ8/DtOE=
 github.com/google/uuid v1.2.0 h1:qJYtXnJRWmpe7m/3XlyhrsLrEURqHRM2kxzoxXqyUDs=
 github.com/google/uuid v1.2.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
-github.com/stretchr/testify v1.3.0 h1:TivCn/peBQ7UY8ooIcPgZFpTNSz0Q2U6UrFlUfqbe0Q=
-github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
+github.com/stretchr/testify v1.7.1 h1:5TQK59W5E3v0r2duFAb7P95B6hEeOyEnHRa8MjYSMTY=
+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=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

Original file line number	Diff line number	Diff line change
`@@ -19,6 +19,9 @@ import (`
`19`	`19`	`"time"`
`20`	`20`
`21`	`21`	`"github.com/google/uuid"`
	`22`	`+ "go.opentelemetry.io/otel"`
	`23`	`+ "go.opentelemetry.io/otel/propagation"`
	`24`	`+ "go.opentelemetry.io/otel/trace"`
`22`	`25`	`)`
`23`	`26`
`24`	`27`	`type key int`
`@@ -374,6 +377,12 @@ func (c *conn) buildRequest(ctx context.Context, query string, params []driver.V`
`374`	`377`	`req.Header.Set("Content-Encoding", "gzip")`
`375`	`378`	`}`
`376`	`379`
	`380`	`+ parentSpan := trace.SpanFromContext(ctx)`
	`381`	`+ if parentSpan.SpanContext().IsValid() {`
	`382`	`+ carrier := propagation.HeaderCarrier(req.Header)`
	`383`	`+ otel.GetTextMapPropagator().Inject(ctx, carrier)`
	`384`	`+ }`
	`385`	`+`
`377`	`386`	`return req, nil`
`378`	`387`	`}`
`379`	`388`