From ff5d3663d7179d1becf51d831eebf5de1ce15cba Mon Sep 17 00:00:00 2001 From: Koptelov Nikita Date: Fri, 17 Jun 2022 14:23:22 +0300 Subject: [PATCH 1/2] WithSpanObserver added --- common.go | 20 +++++++++++++++----- conn.go | 4 ++-- driver.go | 15 +-------------- options.go | 28 ++++++++++++++++++++++++++++ stmt.go | 4 ++-- 5 files changed, 48 insertions(+), 23 deletions(-) create mode 100644 options.go diff --git a/common.go b/common.go index ab64a65..d8784e5 100644 --- a/common.go +++ b/common.go @@ -12,8 +12,14 @@ import ( // ErrUnsupported is an error returned when the underlying driver doesn't provide a given function. var ErrUnsupported = errors.New("operation unsupported by the underlying driver") -// TagQuery is a span tag for SQL queries. -const TagQuery = "query" +// spanTagQuery is a span tag for SQL queries. +var spanTagQuery = "query" + +// SetTagQuery sets a span tag for SQL queries. +// Use in init funcs only. +func SetTagQuery(tag string) { + spanTagQuery = tag +} // SpanNameFunc defines a function which returns a name for the span which is being created on traceable operations. // Passing span naming function is optional, however it gives the user a way to use a custom naming strategy. To allow @@ -23,9 +29,10 @@ type SpanNameFunc func(context.Context) string // tracer defines a set of instances for collecting spans. type tracer struct { - t opentracing.Tracer - nameFunc SpanNameFunc - saveQuery bool + t opentracing.Tracer + nameFunc SpanNameFunc + observerFunc func(context.Context, opentracing.Span) + saveQuery bool } // newSpan creates a new opentracing.Span instance from the given context. @@ -37,6 +44,9 @@ func (t *tracer) newSpan(ctx context.Context) opentracing.Span { opts = append(opts, opentracing.ChildOf(parent.Context())) } span := t.t.StartSpan(name, opts...) + if t.observerFunc != nil { + t.observerFunc(ctx, span) + } return span } diff --git a/conn.go b/conn.go index 7e0b6cd..70a1897 100644 --- a/conn.go +++ b/conn.go @@ -71,7 +71,7 @@ func (c *conn) Exec(query string, args []driver.Value) (driver.Result, error) { func (c *conn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) { s := c.tracer.newSpan(ctx) if c.tracer.saveQuery { - s.SetTag(TagQuery, query) + s.SetTag(spanTagQuery, query) } defer s.Finish() if execerContext, ok := c.conn.(driver.ExecerContext); ok { @@ -107,7 +107,7 @@ func (c *conn) Query(query string, args []driver.Value) (driver.Rows, error) { func (c *conn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (rows driver.Rows, err error) { s := c.tracer.newSpan(ctx) if c.tracer.saveQuery { - s.SetTag(TagQuery, query) + s.SetTag(spanTagQuery, query) } defer s.Finish() if queryerContext, ok := c.conn.(driver.QueryerContext); ok { diff --git a/driver.go b/driver.go index bcbc6ab..324c290 100644 --- a/driver.go +++ b/driver.go @@ -1,6 +1,7 @@ package sql import ( + "context" "database/sql/driver" "github.com/opentracing/opentracing-go" @@ -24,20 +25,6 @@ func NewTracingDriver(d driver.Driver, t opentracing.Tracer, options ...func(*tr return td } -// SpanNameFunction is an option for using a custom span naming function. -func SpanNameFunction(f SpanNameFunc) func(*tracingDriver) { - return func(d *tracingDriver) { - d.tracer.nameFunc = f - } -} - -// SaveQuery is an option for saving SQL queries. -func SaveQuery(f SpanNameFunc) func(*tracingDriver) { - return func(d *tracingDriver) { - d.tracer.saveQuery = true - } -} - // Open implements driver.Driver Open. func (d *tracingDriver) Open(name string) (driver.Conn, error) { c, err := d.driver.Open(name) diff --git a/options.go b/options.go new file mode 100644 index 0000000..072b754 --- /dev/null +++ b/options.go @@ -0,0 +1,28 @@ +package sql + +import ( + "context" + + "github.com/opentracing/opentracing-go" +) + +// SpanNameFunction is an option for using a custom span naming function. +func SpanNameFunction(f SpanNameFunc) func(*tracingDriver) { + return func(d *tracingDriver) { + d.tracer.nameFunc = f + } +} + +// SaveQuery is an option for saving SQL queries. +func SaveQuery(f SpanNameFunc) func(*tracingDriver) { + return func(d *tracingDriver) { + d.tracer.saveQuery = true + } +} + +// WithSpanObserver allows you to modify the span's tags for every span created. +func WithSpanObserver(obsFunc func(context.Context, opentracing.Span)) func(*tracingDriver) { + return func(d *tracingDriver) { + d.tracer.observerFunc = obsFunc + } +} diff --git a/stmt.go b/stmt.go index 9dbeca8..69f05cd 100644 --- a/stmt.go +++ b/stmt.go @@ -35,7 +35,7 @@ func (s *stmt) Query(args []driver.Value) (driver.Rows, error) { func (s *stmt) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) { span := s.tracer.newSpan(ctx) if s.tracer.saveQuery { - span.SetTag(TagQuery, query) + span.SetTag(spanTagQuery, query) } defer span.Finish() if execerContext, ok := s.stmt.(driver.ExecerContext); ok { @@ -52,7 +52,7 @@ func (s *stmt) ExecContext(ctx context.Context, query string, args []driver.Name func (s *stmt) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (rows driver.Rows, err error) { span := s.tracer.newSpan(ctx) if s.tracer.saveQuery { - span.SetTag(TagQuery, query) + span.SetTag(spanTagQuery, query) } defer span.Finish() if queryerContext, ok := s.stmt.(driver.QueryerContext); ok { From 34e2949c54eaa68933361d3a19b6048472b67be9 Mon Sep 17 00:00:00 2001 From: Koptelov Nikita Date: Fri, 17 Jun 2022 14:56:28 +0300 Subject: [PATCH 2/2] build fixed --- driver.go | 1 - 1 file changed, 1 deletion(-) diff --git a/driver.go b/driver.go index 324c290..0417166 100644 --- a/driver.go +++ b/driver.go @@ -1,7 +1,6 @@ package sql import ( - "context" "database/sql/driver" "github.com/opentracing/opentracing-go"