Expose query in the context of the span formatter (#27)

README.md

@@ -197,6 +197,39 @@ func openDB(dsn string) (*sql.DB, error) {
 }
 ```
 
+With traces of `ExecContext()` and `QueryContext()` (either `DB`, `Stmt`, or `Tx`), you could get the SQL query from the context
+using `otelsql.QueryFromContext()`. For example:
+
+```go
+package example
+
+import (
+	"context"
+	"database/sql"
+
+	"github.com/nhatthm/otelsql"
+)
+
+func openDB(dsn string) (*sql.DB, error) {
+	driverName, err := otelsql.Register("my-driver",
+		otelsql.WithSpanNameFormatter(func(ctx context.Context, op string) string {
+			if op != "exec" {
+				return "main-db:" + op
+			}
+
+			query := otelsql.QueryFromContext(ctx)
+
+			// Make span name from the query here and return.
+		}),
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	return sql.Open(driverName, dsn)
+}
+```
+
 [<sub><sup>[table of contents]</sup></sub>](#table-of-contents)
 
 ### Convert Error to Span Status
@@ -442,7 +475,8 @@ func openDB(dsn string) (*sql.DB, error) {
 | `*Result.LastInsertID`  | Disabled. Use `TraceLastInsertID()` to enable |
 | `*Result.RowsAffected`  | Disabled. Use `TraceRowsAffected()` to enable |
 
-`ExecContext`, `QueryContext`, `QueryRowContext`, `PrepareContext` are always traced without query args unless using `TraceQuery()`, `TraceQueryWithArgs()`, or `TraceQueryWithoutArgs()` option.
+`ExecContext`, `QueryContext`, `QueryRowContext`, `PrepareContext` are always traced without query args unless using `TraceQuery()`, `TraceQueryWithArgs()`,
+or `TraceQueryWithoutArgs()` option.
 
 Using `WithDefaultAttributes(...attribute.KeyValue)` will add extra attributes to the recorded spans.
 

context.go

@@ -0,0 +1,20 @@
+package otelsql
+
+import "context"
+
+type queryCtxKey struct{}
+
+// QueryFromContext gets the query from context.
+func QueryFromContext(ctx context.Context) string {
+	query, ok := ctx.Value(queryCtxKey{}).(string)
+	if !ok {
+		return ""
+	}
+
+	return query
+}
+
+// ContextWithQuery attaches the query to the parent context.
+func ContextWithQuery(ctx context.Context, query string) context.Context {
+	return context.WithValue(ctx, queryCtxKey{}, query)
+}

context_test.go

@@ -0,0 +1,23 @@
+package otelsql_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/nhatthm/otelsql"
+)
+
+func TestQueryContext(t *testing.T) {
+	t.Parallel()
+
+	actual := otelsql.QueryFromContext(context.Background())
+	assert.Empty(t, actual)
+
+	ctx := otelsql.ContextWithQuery(context.Background(), "SELECT 1")
+	actual = otelsql.QueryFromContext(ctx)
+	expected := "SELECT 1"
+
+	assert.Equal(t, expected, actual)
+}