database/sqlのドライバをラップしてSQLにコメントをつける

メタデータとしてのSQLコメント

SQLにメタデータのコメントをつけてトレーサビリティをあげる、というのは以前からよく知られたやり方だと思います。
たとえば、Webアプリケーションのユーザー情報取得処理がSELECT name, email FROM usersというクエリをDBに投げていたとして、コントローラーやアクションなどのメタデータのコメントを自動的に追加して、クエリログから実行しているコードを特定したいときなどに有用です。

/* controller: user, action: list */ SELECT name, email FROM users

sqlcommenter

データベースの性能検証をしているときに、Slackに「コメントいれられないかな」と書いたところ、同僚がsqlcommenterというフォーマット仕様とdd-trace-goでの実装について教えてくれました。

https://google.github.io/sqlcommenter/

dd-trace-goの方はDatadogのDatabase Monitoringを有効にしたときに使えるもので(ref)、DBMを有効にしていなかったので機能検証はできなかったのですが、sqlcommenterのGoの実装を簡単に動かしてみました。

package main

import (
	"fmt"

	sqlcommentercore "github.com/google/sqlcommenter/go/core"
	gosql "github.com/google/sqlcommenter/go/database/sql"
	_ "github.com/jackc/pgx/v5/stdlib"
)

func main() {
	dsn := "postgres://[e-mail]:5432/postgres"
	db, err := gosql.Open("pgx", dsn, sqlcommentercore.CommenterOptions{
		Config: sqlcommentercore.CommenterConfig{
			EnableDBDriver:    true,
			EnableApplication: true,
		},
		Tags: sqlcommentercore.StaticTags{Application: "foo"},
	})
	if err != nil {
		panic(err)
	}
	defer db.Close()
	var n any
	for {
		err = db.QueryRow("SELECT pg_sleep(1)").Scan(&n)
		if err != nil {
			panic(err)
		}
		fmt.Printf("n='%+v'\n", n)
	}
}

Performance Insightsやpg_stat_activityでSQLにコメントが追加されていることを確認できます。

database/sqlドライバのラッパーを作る

sqlcommenterのフィールドの定義？を読むとcontrollerやactionなどWebアプリケーションでの利用が想定されているようです。

https://github.com/google/sqlcommenter/blob/871c1c81e014830c3f5f402a9c63d83823132542/go/core/core.go#L29-L40

仕様として任意のキーを使えるかわからなかったのですが、ジョブキューシステムで使うクエリにジョブ名のコメントをつけるなどのユースケースがありそうのなので、SQLにコメントをつけるドライバのラッパーを自分で書いてみました。

package main

import (
	"context"
	"database/sql"
	"database/sql/driver"
	"encoding/json"
	"fmt"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/stdlib"
	_ "github.com/jackc/pgx/v5/stdlib"
)

/////////////////////////////////////////////////////////////////////
// SQLTag
/////////////////////////////////////////////////////////////////////

type ctxKey struct{}
type SQLTag map[string]string

func (t SQLTag) WithContext(ctx context.Context) context.Context {
	return context.WithValue(ctx, ctxKey{}, t)
}

func SQLTagFromContext(ctx context.Context) SQLTag {
	value := ctx.Value(ctxKey{})
	if value == nil {
		return nil
	}
	return value.(SQLTag)
}

/////////////////////////////////////////////////////////////////////
// Conn
/////////////////////////////////////////////////////////////////////

var (
	_ driver.Conn               = (*MyConn)(nil)
	_ driver.QueryerContext     = (*MyConn)(nil)
	_ driver.ExecerContext      = (*MyConn)(nil)
	_ driver.ConnPrepareContext = (*MyConn)(nil)
)

type MyConn struct {
	driver.Conn
}

func (c *MyConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
	queryer, ok := c.Conn.(driver.QueryerContext)
	if !ok {
		return nil, driver.ErrSkip
	}
	tag := SQLTagFromContext(ctx)
	if tag != nil {
		js, _ := json.Marshal(tag)
		query = fmt.Sprintf("/* %s */ %s", js, query)
	}
	return queryer.QueryContext(ctx, query, args)
}

func (c *MyConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
	execer, ok := c.Conn.(driver.ExecerContext)
	// 省略
}

func (c *MyConn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
	preparer, ok := c.Conn.(driver.ConnPrepareContext)
	// 省略
}

/////////////////////////////////////////////////////////////////////
// Driver
/////////////////////////////////////////////////////////////////////

var (
	_ driver.Driver        = (*MyDriver)(nil)
	_ driver.DriverContext = (*MyDriver)(nil)
)

type MyDriver struct {
	wrapped driver.Driver
}

func (d *MyDriver) Open(name string) (driver.Conn, error) {
	conn, err := d.wrapped.Open(name)
	if err != nil {
		return nil, err
	}
	return &MyConn{conn}, nil
}

func (d *MyDriver) OpenConnector(name string) (driver.Connector, error) {
	connector, err := d.wrapped.(driver.DriverContext).OpenConnector(name)
	if err != nil {
		return nil, err
	}
	return &MyConnector{connector}, nil
}

/////////////////////////////////////////////////////////////////////
// Connector
/////////////////////////////////////////////////////////////////////

var (
	_ driver.Connector = (*MyConnector)(nil)
)

type MyConnector struct {
	wrapped driver.Connector
}

func (c *MyConnector) Connect(ctx context.Context) (driver.Conn, error) {
	conn, err := c.wrapped.Connect(ctx)
	if err != nil {
		return nil, err
	}
	return &MyConn{conn}, nil
}

func (c *MyConnector) Driver() driver.Driver {
	return &MyDriver{c.wrapped.Driver()}
}

/////////////////////////////////////////////////////////////////////
// main
/////////////////////////////////////////////////////////////////////

func main() {
	dsn := "postgres://[e-mail]:5432/postgres"
	connCfg, err := pgx.ParseConfig(dsn)
	if err != nil {
		panic(err)
	}
	connector := &MyConnector{stdlib.GetConnector(*connCfg)}
	db := sql.OpenDB(connector)
	defer db.Close()
	var n any
	for {
		ctx := context.Background()
		tag := SQLTag{"foo": "bar"}
		ctx = tag.WithContext(ctx)
		err = db.QueryRowContext(ctx, "SELECT pg_sleep(1)").Scan(&n)
		if err != nil {
			panic(err)
		}
		fmt.Printf("n='%+v'\n", n)
	}
}

contextにセットした任意のキーバリューをJSON^[1]に変換して、SQLにコメントできるようになりました。

脚注

1. */がエスケープされないのでJSONエンコードする方法を実際に使うのはやめた方がよさそうです