DataDog/dd-trace-goがgorm.ioをサポートしている話

はじめに

さくっとした投稿です。
Gormメジャーバージョンアップがあり、importパスがgithub.com/jinzhu/gormからgorm.io/gormになったのはご存知の方も多いと思います。

DatadogのAPMを利用していた場合、gormとDatadogのトレースを紐づけるライブラリhttps://github.com/DataDog/dd-trace-go/ の開発が追いついておらず、gorm.io/gormに対応できていなくて困っていた方も多いかと思います。

ふとレポジトリをあさっていたら、今までの
gopkg.in/DataDog/dd-trace-go.v1/contrib/jinzhu/gorm
ではなく、
gopkg.in/DataDog/dd-trace-go.v1/contrib/gorm.io/gorm.v1
に、3か月前の2021/1/5に爆誕しているのを発見しました。

こいつに触れてあげている情報がネット上だと皆無だったので、気づいてあげられるようブログにすることにしました。

どう対応すればいいか

As Is

今までは以下のようにしてました。

import (
  .
  .
  gormtrace "gopkg.in/DataDog/dd-trace-go.v1/contrib/jinzhu/gorm"
)
// アプリ起動時に一度だけ実行される
func Connection() *gorm.DB {
  .
  .      
  db, _ := gorm.Open("postgres", dsn)
  db = gormtrace.WithCallbacks(db)
  db.DB()
  db.DB().SetConnMaxLifetime(MaxLifetime)
  db.DB().SetMaxIdleConns(MaxIdle)
  db.DB().SetMaxOpenConns(MaxOpen)
  db.LogMode(Debug)
  return db
}

import (
  .
  .
  "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
)

// sql実行時
// txはConnectionの戻り値dbに対して、db.Beginしたもの
func ExecQuery(ctx context.Context, tx *gorm.DB) {
  tx = gormtrace.WithContext(ctx, tx)
  if result := tx.Save(&record); result.Error != nil {
  ..

To Be

import (
  .
  .
  gormtrace "gopkg.in/DataDog/dd-trace-go.v1/contrib/gorm.io/gorm.v1"
)

// アプリ起動時に一度だけ実行される
func Connection() *gorm.DB {
  // gormtraceはOpenにのみ必要
  db, _ := gormtrace.Open(postgres.New(postgres.Config{
       DriverName: "postgres",
       DSN:        dsn,
    }), &gorm.Config{
       Logger: logger.Default.LogMode(LogLevel),
    }, gormtrace.WithAnalytics(true),
  )
  db, err := db.DB() 
  db.SetConnMaxLifetime(MaxLifetime)
  db.SetMaxIdleConns(MaxIdle)
  db.SetMaxOpenConns(MaxOpen)

// sql実行時
// txはConnectionの戻り値dbに対して、db.Beginしたもの
func ExecQuery(ctx context.Context,tx *gorm.DB) {
  // tx.WithContext(ctx)をすると、tx.Statement.Contextに格納
  // traceするときにtx.Statement.Contextをみてくれている。
  if result := tx.WithContext(ctx).Save(&record)result.Error != nil {
  ..

少しシンプルになっていることがわかる。
これはgormがBegin意外にもcontextを与えられるようになったからだろう。

ソースでいうとこの辺り

  // db.Statement.Contextから取得
  ctx := db.Statement.Context
  // クエリ実行前のcallbackにて値が格納されている
  t, ok := ctx.Value(gormSpanStartTimeKey).(time.Time) 
  if !ok {
    return
  }
  opts := []ddtrace.StartSpanOption{
    tracer.StartTime(t),
    tracer.ServiceName(cfg.serviceName),
    tracer.SpanType(ext.SpanTypeSQL),
    tracer.ResourceName(db.Statement.SQL.String()),
  }
  if !math.IsNaN(cfg.analyticsRate) {
    opts = append(opts, tracer.Tag(ext.EventSampleRate, cfg.analyticsRate))
  }
  // ctxからspanを取得可能
  span, _ := tracer.StartSpanFromContext(ctx, operationName, opts...)
  span.Finish(tracer.WithError(db.Error))
}