[Tech Blog] Writing simple tracing wrappers for golang pgx and http.Client for Sentry

Hello! This is Vasily, a backend engineer in AnyLogi,
In this article, I am going to show you some examples of custom tracing instrumentation for sentry, that we use in our go projects.

Imagine a situation when you want to instrument some code for a well-known library, and you think, oh, it is a typical case, someone must have solved this before and published a solution. But it turns out there is no such thing, all you can find is outdated GitHub issues and vague answers on how it can be implemented. It can be discouraging, and you might stop right there and give up on instrumentation because you don’t want to invent a bicycle or simply don’t know how to do it. In my humble opinion, it is not an excuse not to trace something, especially if it is a critical component. After all, no one would argue that tracing is so important. So here are two simple examples showing how it can be done.

Note: Although the examples use Sentry, it can be easily rewritten to any other tool

Instrumenting https://github.com/jackc/pgx (4.6k ⭐)

pgx is a pure Go driver and toolkit for PostgreSQL.

We use this library in AnyLogi Go projects because it is fast, gives full control, supports all postgres data types.

However, surprisingly it doesn’t have written wrappers for tracing (or at least I wasn’t able to find them).

Also, it doesn’t have any hooks or middlewares, where we could cram our tracing.

Nonetheless, we have two options:

Option 1: Wrapping of pgx.Conn (pgxpool.Conn)

A straightforward solution by wrapping the pgx connection interface

The list of interfaces:

// implementations: pgx.Conn, pgxpool.Conn
// not a pgx interface 
type QuerierConn interface {
	pgxtype.Querier
	Conn
}
// not a pgx interface
type Conn interface {
	Begin(context.Context) (pgx.Tx, error)
	BeginTx(context.Context, pgx.TxOptions) (pgx.Tx, error)
}
// pgx interface
type Querier interface {
	Exec(ctx context.Context, sql string, arguments ...interface{}) (pgconn.CommandTag, error)
	Query(ctx context.Context, sql string, optionsAndArgs ...interface{}) (pgx.Rows, error)
	QueryRow(ctx context.Context, sql string, optionsAndArgs ...interface{}) pgx.Row
}
// pgx interface
type Tx interface {
	Querier
	Commit(ctx context.Context) error
	Rollback(ctx context.Context) error
	...
}

Notice that the connection (QuerierConn) and the transaction (Tx) are both able to execute a query (Querier interface), which makes this part a little bit tricky as we need to instrument Querier in both QuerierConn and Tx and avoid code duplication.

The wrapper code:

pgx_tracing_wrapper.go

type pgxQuerierConnWrapper struct {
	Conn
	pgxtype.Querier
}
// pgxQuerierConnWrapper uses pgxQuerierWrapper
// the same wrapper will be used again to partially wrap pgx.Tx 
func NewPgxQuerierConnWrapper(querierConn postgres.QuerierConn) *pgxQuerierConnWrapper {
	return &pgxQuerierConnWrapper{
		Conn:    querierConn,
		Querier: newPgxQuerierWrapper(querierConn),
	}
}

func (p pgxQuerierConnWrapper) Begin(ctx context.Context) (pgx.Tx, error) {
	tx, err := p.Conn.Begin(ctx)
	return newPgxTxWrapper(tx), err
}

func (p pgxQuerierConnWrapper) BeginTx(ctx context.Context, options pgx.TxOptions) (pgx.Tx, error) {
	tx, err := p.Conn.BeginTx(ctx, options)
	return newPgxTxWrapper(tx), err
}

type pgxTxWrapper struct {
	pgx.Tx
	pgxtype.Querier
}
// pgxTxWrapper embeds two interfaces and as you can see both of them are derived from pgx.Tx,
// pgxtype.Querier is wrapped by pgxQuerierWrapper
// and pgx.Tx covers the rest methods
// in result we have pgx.Tx partially wrapped and other methods remained untouched
func newPgxTxWrapper(tx pgx.Tx) *pgxTxWrapper {
	return &pgxTxWrapper{
		Tx:      tx,
		Querier: newPgxQuerierWrapper(tx),
	}
}

func (p *pgxTxWrapper) Exec(ctx context.Context, sql string, arguments ...interface{}) (pgconn.CommandTag, error) {
	return p.Querier.Exec(ctx, sql, arguments...)
}

func (p *pgxTxWrapper) Query(ctx context.Context, sql string, optionsAndArgs ...interface{}) (pgx.Rows, error) {
	return p.Querier.Query(ctx, sql, optionsAndArgs...)
}

func (p *pgxTxWrapper) QueryRow(ctx context.Context, sql string, optionsAndArgs ...interface{}) pgx.Row {
	return p.Querier.QueryRow(ctx, sql, optionsAndArgs...)
}

// pgxtype.Querier wrapper
type pgxQuerierWrapper struct {
	pgxtype.Querier
}

func newPgxQuerierWrapper(querier pgxtype.Querier) *pgxQuerierWrapper {
	return &pgxQuerierWrapper{Querier: querier}
}

func (p *pgxQuerierWrapper) Exec(ctx context.Context, sql string, arguments ...interface{}) (pgconn.CommandTag, error) {
	span := sentry.StartSpan(ctx, "postgres Exec")
	defer span.Finish()
	if span.Data == nil {
		span.Data = make(map[string]interface{})
	}
	span.Data["sql"] = sql
	return p.Querier.Exec(ctx, sql, arguments...)
}

func (p *pgxQuerierWrapper) Query(ctx context.Context, sql string, optionsAndArgs ...interface{}) (pgx.Rows, error) {
	span := sentry.StartSpan(ctx, "postgres Query")
	defer span.Finish()
	if span.Data == nil {
		span.Data = make(map[string]interface{})
	}
	span.Data["sql"] = sql
	return p.Querier.Query(ctx, sql, optionsAndArgs...)
}

func (p *pgxQuerierWrapper) QueryRow(ctx context.Context, sql string, optionsAndArgs ...interface{}) pgx.Row {
	span := sentry.StartSpan(ctx, "postgres QueryRow")
	defer span.Finish()
	if span.Data == nil {
		span.Data = make(map[string]interface{})
	}
	span.Data["sql"] = sql
	return p.Querier.QueryRow(ctx, sql, optionsAndArgs...)
}

Example:

var conn QuerierConn

conn, err := pgxpool.Connect(context.Background(), os.Getenv("DATABASE_URL"))
if err != nil {
	log.Fatal(err)
}
conn := NewPgxQuerierConnWrapper(conn)

// pass wrapped implementation of the interface
orderRepository := NewOrderRepository(conn)
...

this is fairly simplified code, omitted some parts like Commit or Rollback on purpose, as it was becoming too verbose, if you think so as well, there is a second option

Option 2: Using pgx connection logger for tracing

pgx allows injecting a custom Logger exposing a plain interface and even provides some adapters for popular loggers like logrus or zap. Obviously, a logger shouldn’t be used for tracing as this violates the S letter in SOLID, but it turned out it is the only proper place with context propagation to perform tracing. Even the pgx maintainer himself suggested using this hacky solution. https://github.com/jackc/pgx/issues/814#issuecomment-682617143

Here is the code, it is way shorter than the 1st option

// TracingLogger does two things at once, logging and tracing
type TracingLogger struct {
	logger *logrusadapter.Logger
}

func NewTracingLogger(logger logrus.FieldLogger) *TracingLogger {
	return &TracingLogger{logger: logrusadapter.NewLogger(logger)}
}

func (l *TracingLogger) Log(ctx context.Context, level pgx.LogLevel, msg string, data map[string]interface{}) {
	l.logger.Log(ctx, level, msg, data)

	span := sentry.StartSpan(ctx, "postgres "+msg)
	defer span.Finish()

	if data != nil {
		if span.Data == nil {
			span.Data = make(map[string]interface{})
		}
		for key, value := range data {
			value, ok := value.(string)
			if ok {
				span.Data[key] = value
			}
		}
		// since this method is called after a query is done, 
		// we need to correct span start time 
		execTime, ok := data["time"].(time.Duration)
		if ok {
			span.StartTime = time.Now().Add(-execTime)
		}
	}
}

Usage:

config, err := pgxpool.ParseConfig(os.Getenv("DATABASE_URL"))
if err != nil {
    log.Fatal(err)
}
config.ConnConfig.Logger = NewTracingLogger(logger)

conn, err = pgxpool.ConnectConfig(ctx, config)

A sentry span visualization in the dashboard:

Instrumenting https://pkg.go.dev/net/http#Client

http.Client – Go standard library HTTP client implementation

Unfortunately Sentry sdk doesn’t have a wrapper for this crucial component, but no worries, it is not hard to implement.

One of the easiest ways to instrument http.Client is by wrapping http.Transport, it can be implemented as follows:

type tracingTransport struct {
	http.RoundTripper
}

func NewTracingTransport(roundTripper http.RoundTripper) *tracingTransport {
	return &tracingTransport{RoundTripper: roundTripper}
}

func (t *tracingTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	operationName := fmt.Sprintf("HTTP %s %s", req.Method, req.URL.String())
	span := sentry.StartSpan(req.Context(), operationName)
	defer span.Finish()

	span.SetTag("url", req.URL.String())
	if span.Data == nil {
		span.Data = make(map[string]interface{})
	}
	// reading body from the request body and fill it again
	body, err := ioutil.ReadAll(req.Body)
	if err != nil {
		return nil, err
	}
	// Be careful with including sensitive information in the span,
	// request body and response may have private user data, which we wouldn't want to expose,
	// authorization header also is a good example of sensitive data.
	span.Data["body"] = string(body)

	req.Body = ioutil.NopCloser(bytes.NewBuffer(body))
	// adding sentry header for distributed tracing
	req.Header.Add("sentry-trace", span.TraceID.String())

	response, err := t.RoundTripper.RoundTrip(req)

	span.Data["http_code"] = response.StatusCode
	// could additionally add the response to the span data

	return response, err
}

Usage:

client := &http.Client{
    Transport: NewTracingTransport(http.DefaultTransport),
}

Result:

Conclusion:

Most popular libs and tracing tools provide SDKs to avoid all of this, but sometimes there is no other way than do it by yourself. When it happens, it is helpful to be ready to dive into documentation/source code and produce a viable solution. In practice, there are always ways to instrument code, some of them are pretty, some are not, just need to find them. These examples showcased simple techniques with wrapping interfaces and one with a nonstandard approach. Hope someone will find it helpful.

Thank you.