[OpenTelemetry 検証7] pgx.QueryTracer で PostgreSQL のクエリを自動計装してみる

はじめに

[OpenTelemetry 検証6] Grafana ダッシュボードを作成・編集してみる の続きです。

前回は Grafana ダッシュボードをプロビジョニングし、HTTP レイテンシとリクエストレートを 1 画面で確認できる状態を作りましたが、今回は pgx の QueryTracer インターフェースを実装し、PostgreSQL への SQL クエリを span として記録できる仕組みを作っていきたいと思います。

---

pgx.QueryTracer について

pgx は Go で広く使われる PostgreSQL ドライバで、クエリの実行前後にフックできる pgx.QueryTracer インターフェースを提供しています。

type QueryTracer interface {
    TraceQueryStart(ctx context.Context, conn *Conn, data TraceQueryStartData) context.Context
    TraceQueryEnd(ctx context.Context, conn *Conn, data TraceQueryEndData)
}

このインターフェースを実装した構造体を *pgxpool.Config に設定するだけで、そのコネクションプール経由で実行される全ての SQL クエリが自動的にトレーシングの対象になります。

※なので repository 層の各メソッドに個別のトレース処理を書く必要もありません

参考: pgx/tracer.go

---

pgxTracer を実装してみる

早速試してみたいと思うので、今回は pkg/postgresql/driver.go に pgxTracer を実装しました。

ざっとこのような実装にしています。

var tracer = otel.Tracer("postgresql")

type querySpanKey struct{}

type pgxTracer struct{}

// TraceQueryStart は SQL クエリの実行開始時にスパンを開始し、db.system / db.statement 属性を付与する。
// pgx がクエリを実行するたびに自動的に呼び出す仕組みになっている。
func (t *pgxTracer) TraceQueryStart(ctx context.Context, _ *pgx.Conn, data pgx.TraceQueryStartData) context.Context {
    ctx, span := tracer.Start(ctx, "pgx.Query",
        trace.WithSpanKind(trace.SpanKindClient),
        trace.WithAttributes(
            attribute.String("db.system", "postgresql"),
            attribute.String("db.statement", data.SQL),
        ),
    )
    return context.WithValue(ctx, querySpanKey{}, span)
}

// TraceQueryEnd は SQL クエリの実行終了時にスパンを終了する。
// エラーがあれば記録した上でスパンを閉じる。
func (t *pgxTracer) TraceQueryEnd(ctx context.Context, _ *pgx.Conn, data pgx.TraceQueryEndData) {
    span, ok := ctx.Value(querySpanKey{}).(trace.Span)
    if !ok {
        return
    }
    if data.Err != nil {
        span.RecordError(data.Err)
        span.SetStatus(codes.Error, data.Err.Error())
    }
    span.End()
}

ポイントとしては、TraceQueryStart と TraceQueryEnd は、pgx SDK が内部でこのインターフェースのメソッドを呼び出す仕組みになっているため、アプリケーションコードから直接呼び出されることはないことです。

querySpanKey{} はコンテキストのキーとして使う privateな構造体で、他パッケージの key と被らないようにしています。

---

pgxpool への注入

NewPool 関数で pgxpool を生成する際に、*pgxpool.Config (ここではconfig.ConnConfig.Tracer) に pgxTracer を設定しています。

このように設定することで、このプール経由の全 SQL クエリがスパンとして記録されるようになります。

func NewPool(ctx context.Context, dsn string) (*pgxpool.Pool, error) {
    config, err := pgxpool.ParseConfig(dsn)
    if err != nil {
        return nil, err
    }
    config.ConnConfig.Tracer = &pgxTracer{} // 全クエリを自動トレーシングの対象にする
    return pgxpool.NewWithConfig(ctx, config)
}

---

スパンの階層構造

今回の実装により、Grafana (Tempo) で確認できるスパンの階層は次のようになります。

http-server
└── ArticleUsecase.GetByID / Create      (SpanKindInternal)
    └── ArticleRepository.FindByID / Create  (SpanKindClient)
        └── pgx.Query                        (SpanKindClient, 自動計装) // これが新しく追加された

repository 層の手動スパンと pgxTracer による自動スパンが入れ子になり、「どの usecase が」「どの repository メソッドを呼んで」「どんな SQL を発行したか」が 1 つのトレースでで分かるようになります。

---

Grafana で確認してみる

pgx.Query の span には db.statement 属性として実際の SQL が記録されているため、どのクエリに時間がかかっているかが確認できるようになります。

まずはざっとこんな感じで適当にリクエストを投げてみます。

この状態で Trace の画面にいくと以下のようになっていることが確認できます。

今回追加した pgx.Query が Trace に紐づけられていることが確認できます。

また、db.statement を見ると以下のようなクエリが実行されていることもわかります。

SELECT
	id,
	title,
	body,
	status,
	created_at,
	updated_at
FROM articles
WHERE id = $1

今回は1クエリに依存した trace になっていますが、少し規模が大きいシステムでは複数のクエリが1つのリクエスト内で実行されることが多いかと思います。

この時、どのクエリで時間を要しているのかを特定する際に この trace が役立つことが多いため、運用担当者だけでなく開発者側もこの辺りの知見を抑えておくと良さそうです。

---

おわりに

今回は pgx.QueryTracer インターフェースを実装し、PostgreSQL への SQL クエリの自動計装を実装してみました。

アプリケーション側のコードを変更せずに、わずか数行で設定できるのはとても便利に思えます。

また、手動 span と組み合わせることで、HTTP リクエストから SQL 実行までの処理フローを 1 つの trace として可視化できる点も強力です。

次回は otelhttp.NewTransport を使った HTTP クライアントの自動計装で、複数のサーバー間で traceparent ヘッダーを自動付与し、いわゆるマイクロサービス間のトレース伝播の動作などを確認してみたいと思います。