goのdriver.DriverでAWS X-RayによるTracingとSQL Commentを両立させる

aws-xray-sdk-go の xray.SQLContext を利用することで、 AWS X-Ray のtracingにSQLを発行した際のセグメントを載せることができます。

発行したSQL文はセグメントの SQL タブから確認することができます。ただもうちょっと欲張って、実行箇所をSQL Commentとして埋め込むことでセグメントを見れば発行箇所がわかるようにしたいと考えました。

Go言語におけるSQL Commentの埋め込み方はSongmuさんの以下のブログが詳しいです。

songmu.jp

まずはSQL Commentの埋め込み部分について、以下のように実装してみました。

package myapp

import (
    "context"
    "database/sql"
    "database/sql/driver"
    "fmt"
    "runtime"
    "strings"
    "sync"
    "time"

    "github.com/aws/aws-xray-sdk-go/xray"
    proxy "github.com/shogo82148/go-sql-proxy"

    // for using mysql driver
    _ "github.com/go-sql-driver/mysql"
)

type proxyHooksManager []func(ctx context.Context, stmt *proxy.Stmt, args []driver.NamedValue)

func (hooks *proxyHooksManager) Add(fn func(ctx context.Context, stmt *proxy.Stmt, args []driver.NamedValue)) {
    *hooks = append(*hooks, fn)
}

func (hooks proxyHooksManager) Run(ctx context.Context, stmt *proxy.Stmt, args []driver.NamedValue) (interface{}, error) {
    if stmt.Stmt != nil {
        return nil, nil
    }
    for _, fn := range hooks {
        fn(ctx, stmt, args)
    }
    return nil, nil
}

func (hooks proxyHooksManager) PrePrepare(ctx context.Context, stmt *proxy.Stmt) (interface{}, error) {
    return hooks.Run(ctx, stmt, nil)
}

func (hooks proxyHooksManager) PreExecAndQuery(ctx context.Context, stmt *proxy.Stmt, args []driver.NamedValue) (interface{}, error) {
    return hooks.Run(ctx, stmt, args)
}

var (
    proxyHooks  proxyHooksManager
)

func init() {
    const (
        ignorePath = "/path/to/myapp"
    )
    proxyHooks.Add(func(ctx context.Context, stmt *proxy.Stmt, args []driver.NamedValue) {
        var (
            pc   uintptr
            file string
            line int
        )
        // 適宜調整してください
        for i := 1; i < 100; i++ {
            pc, file, line, _ = runtime.Caller(i)
            if !strings.Contains(file, ignorePath) {
                break
            }
        }
        fn := runtime.FuncForPC(pc)
        stmt.QueryString = fmt.Sprintf("/* %s (%s:%d) */ %s", fn.Name(), file, line, stmt.QueryString)
    })
}

proxyHooks という変数を用意したのは、SQL Commentを埋め込むのと同じタイミングで、テストからも何かしらの処理の注入して実行できるようにしたかったからです。 aws-xray-sdk-go では attrHook という変数でテスト時の処理の注入を可能にしていて、それでも特に問題はないですが、sliceで保持しておけばテスト用の挙動かどうかという境界は無くせるかな、と。

あとは sql.Open を行う際に1度だけ実行するコードを定義します

// これを事前に定義しておく
var driverSetup sync.Once

func Open(dsn string) {
    driverSetup.Do(func() {
        // わざとOpenして `*sql.DB` のオブジェクトを取得する
        db, err := xray.SQLContext("mysql", "")
        if err != nil {
            panic(err)
        }
        defer db.Close()
        sql.Register("mysql:myapp", proxy.NewProxyContext(db.Driver(), &proxy.HooksContext{
            PrePrepare: proxyHooks.PrePrepare,
            PreExec:    proxyHooks.PreExecAndQuery,
            PreQuery:   proxyHooks.PreExecAndQuery,
        }))
    })

    db, err := sql.Open("mysql:myapp", dsn)
    if err != nil {
        panic(err)
    }

    // snip...
}

キモはDSNを指定せずに、わざと xray.SQLContext を実行して *sql.DB のオブジェクトを取得していることで、 ここで得られる driver.Driver を proxy.NewProxyContext の第一引数に指定しています。 そしてアプリケーションが利用する *sql.DB を取得する際は、 proxy.NewProxyContext で指定したドライバー名を使用します。

アプリケーション側が使用する *sql.DB が持っている driver.Driver は、こんな入れ子構造のオブジェクトを利用する形になります。 これにより、最上段の go-sql-proxy によるドライバではSQL Commentを埋め込み、処理を移譲した先のX-Rayのドライバではtracingを行い、更にそこから移譲した先のmysqlのドライバで実際に接続先のDBに向けてクエリを発行する、という流れになります。

今回はX-RayとSQL Commentの組み合わせでしたが、ドライバの移譲を重ねることで機能を足すことができるので、New Relicによる nrmysql 等でも同様のアプローチでが可能だと考えています。