Add SQL Commenter (#112)

* Add SQL commenter

* Apply SQL commenter to all query methods

* Fix tests

* Fix docs

* Update CHANGELOG

* Clean comments

* Add experimental notice to CHANGELOG

CHANGELOG.md

@@ -8,6 +8,12 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ## [Unreleased]
 
+### Added
+
+- `WithSQLCommenter` option to enable context propagation for database by injecting a comment into SQL statements. (#112)
+
+  This is an experimental feature and may be changed or removed in a later release.
+
 ## [0.15.0] - 2022-07-11
 
 ### ⚠️ Notice ⚠️

commenter.go

@@ -0,0 +1,67 @@
+// Copyright Sam Xie
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package otelsql
+
+import (
+	"context"
+	"fmt"
+	"net/url"
+	"strings"
+
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/propagation"
+)
+
+type commentCarrier []string
+
+var _ propagation.TextMapCarrier = (*commentCarrier)(nil)
+
+func (c *commentCarrier) Keys() []string { return nil }
+
+func (c *commentCarrier) Get(string) string { return "" }
+
+func (c *commentCarrier) Set(key, value string) {
+	*c = append(*c, fmt.Sprintf("%s='%s'", url.QueryEscape(key), url.QueryEscape(value)))
+}
+
+func (c *commentCarrier) Marshal() string {
+	return strings.Join(*c, ",")
+}
+
+type commenter struct {
+	enabled    bool
+	propagator propagation.TextMapPropagator
+}
+
+func newCommenter(enabled bool) *commenter {
+	return &commenter{
+		enabled:    enabled,
+		propagator: otel.GetTextMapPropagator(),
+	}
+}
+
+func (c *commenter) withComment(ctx context.Context, query string) string {
+	if !c.enabled {
+		return query
+	}
+
+	var cc commentCarrier
+	c.propagator.Inject(ctx, &cc)
+
+	if len(cc) == 0 {
+		return query
+	}
+	return fmt.Sprintf("%s /*%s*/", query, cc.Marshal())
+}

commenter_test.go

@@ -0,0 +1,85 @@
+// Copyright Sam Xie
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package otelsql
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/otel/baggage"
+	"go.opentelemetry.io/otel/propagation"
+	"go.opentelemetry.io/otel/trace"
+)
+
+func TestCommenter_WithComment(t *testing.T) {
+	query := "foo"
+
+	traceID, err := trace.TraceIDFromHex("a3d3b88cf7994e554c1afbdceec1620b")
+	require.NoError(t, err)
+	spanID, err := trace.SpanIDFromHex("683ec6a9a3a265fb")
+	require.NoError(t, err)
+	traceState, err := trace.ParseTraceState("rojo=00f067aa0ba902b7,congo=t61rcWkgMzE")
+	require.NoError(t, err)
+	ctx := trace.ContextWithSpanContext(context.Background(), trace.NewSpanContext(trace.SpanContextConfig{
+		TraceID:    traceID,
+		SpanID:     spanID,
+		TraceFlags: 0x1,
+		TraceState: traceState,
+	}))
+
+	m1, err := baggage.NewMember("foo", "bar")
+	require.NoError(t, err)
+	b, err := baggage.New(m1)
+	require.NoError(t, err)
+	ctx = baggage.ContextWithBaggage(ctx, b)
+
+	testCases := []struct {
+		name     string
+		enabled  bool
+		ctx      context.Context
+		expected string
+	}{
+		{
+			name:     "empty context",
+			enabled:  true,
+			ctx:      context.Background(),
+			expected: query,
+		},
+		{
+			name:     "context with disable",
+			enabled:  false,
+			ctx:      ctx,
+			expected: query,
+		},
+		{
+			name:     "context",
+			enabled:  true,
+			ctx:      ctx,
+			expected: query + " /*tracestate='rojo%3D00f067aa0ba902b7%2Ccongo%3Dt61rcWkgMzE',traceparent='00-a3d3b88cf7994e554c1afbdceec1620b-683ec6a9a3a265fb-01',baggage='foo%3Dbar'*/",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			c := newCommenter(tc.enabled)
+			c.propagator = propagation.NewCompositeTextMapPropagator(propagation.TraceContext{}, propagation.Baggage{})
+
+			result := c.withComment(tc.ctx, query)
+			assert.Equal(t, tc.expected, result)
+		})
+	}
+}

config.go

@@ -57,6 +57,16 @@ type config struct {
 	// SpanNameFormatter will be called to produce span's name.
 	// Default use method as span name
 	SpanNameFormatter SpanNameFormatter
+
+	// SQLCommenterEnabled enables context propagation for database
+	// by injecting a comment into SQL statements.
+	//
+	// Experimental
+	//
+	// Notice: This config is EXPERIMENTAL and may be changed or removed in a
+	// later release.
+	SQLCommenterEnabled bool
+	SQLCommenter        *commenter
 }
 
 // SpanOptions holds configuration of tracing span to decide
@@ -127,6 +137,8 @@ func newConfig(options ...Option) config {
 		metric.WithInstrumentationVersion(Version()),
 	)
 
+	cfg.SQLCommenter = newCommenter(cfg.SQLCommenterEnabled)
+
 	var err error
 	if cfg.Instruments, err = newInstruments(cfg.Meter); err != nil {
 		otel.Handle(err)

config_test.go

@@ -46,6 +46,7 @@ func TestNewConfig(t *testing.T) {
 			semconv.DBSystemMySQL,
 		},
 		SpanNameFormatter: &defaultSpanNameFormatter{},
+		SQLCommenter:      newCommenter(false),
 	}, cfg)
 	assert.NotNil(t, cfg.Instruments)
 }

conn.go

@@ -103,7 +103,7 @@ func (c *otConn) ExecContext(ctx context.Context, query string, args []driver.Na
 	)
 	defer span.End()
 
-	res, err = execer.ExecContext(ctx, query, args)
+	res, err = execer.ExecContext(ctx, c.cfg.SQLCommenter.withComment(ctx, query), args)
 	if err != nil {
 		recordSpanError(span, c.cfg.SpanOptions, err)
 		return nil, err
@@ -141,7 +141,7 @@ func (c *otConn) QueryContext(ctx context.Context, query string, args []driver.N
 		defer span.End()
 	}
 
-	rows, err = queryer.QueryContext(queryCtx, query, args)
+	rows, err = queryer.QueryContext(queryCtx, c.cfg.SQLCommenter.withComment(queryCtx, query), args)
 	if err != nil {
 		recordSpanError(span, c.cfg.SpanOptions, err)
 		return nil, err
@@ -170,7 +170,7 @@ func (c *otConn) PrepareContext(ctx context.Context, query string) (stmt driver.
 		defer span.End()
 	}
 
-	stmt, err = preparer.PrepareContext(ctx, query)
+	stmt, err = preparer.PrepareContext(ctx, c.cfg.SQLCommenter.withComment(ctx, query))
 	if err != nil {
 		recordSpanError(span, c.cfg.SpanOptions, err)
 		return nil, err

option.go

@@ -72,3 +72,23 @@ func WithMeterProvider(provider metric.MeterProvider) Option {
 		cfg.MeterProvider = provider
 	})
 }
+
+// WithSQLCommenter will enable or disable context propagation for database
+// by injecting a comment into SQL statements.
+//
+// e.g., a SQL query
+//  SELECT * from FOO
+// will become
+//  SELECT * from FOO /*traceparent='00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01',tracestate='congo%3Dt61rcWkgMzE%2Crojo%3D00f067aa0ba902b7'*/
+//
+// This option defaults to disable.
+//
+// Experimental
+//
+// Notice: This option is EXPERIMENTAL and may be changed or removed in a
+// later release.
+func WithSQLCommenter(enabled bool) Option {
+	return OptionFunc(func(cfg *config) {
+		cfg.SQLCommenterEnabled = enabled
+	})
+}

option_test.go

@@ -63,6 +63,11 @@ func TestOptions(t *testing.T) {
 			option:         WithMeterProvider(meterProvider),
 			expectedConfig: config{MeterProvider: meterProvider},
 		},
+		{
+			name:           "WithSQLCommenter",
+			option:         WithSQLCommenter(true),
+			expectedConfig: config{SQLCommenterEnabled: true},
+		},
 	}
 
 	for _, tc := range testCases {

utils_test.go

@@ -141,6 +141,7 @@ func newMockConfig(t *testing.T, tracer trace.Tracer) config {
 		Instruments:       instruments,
 		Attributes:        []attribute.KeyValue{defaultattribute},
 		SpanNameFormatter: &defaultSpanNameFormatter{},
+		SQLCommenter:      newCommenter(false),
 	}
 }