使用OpenCensus跟踪Gorm查询

Gorm作为Go语言中很常用的一个ORM库，功能非常强大。应用程序的大量时间都花在通过gorm与数据库连接上面，所以我们想在链路跟踪中获得更好的视图。

幸运的是，Gorm有完美的钩子，我们可以通过Callbacks API将跟踪功能注入到数据库处理当中。Callbacks API允许我们为Gorm提供在查询生命周期的特定部分中执行相应的函数，或者允许您在传统的中间件方法中更改查询行为，或者在我们的例子中，为可观察性提取数据。

func beforeQuery(scope *gorm.DB) {
    // do stuff!
}

db.Callback().
    Create().
    Before("gorm:query").
    Register("instrumentation:before_query", beforeQuery)

这篇文章的目标是在我们的Gorm查询中引入链路跟踪，为了做到这一点，我们需要同时捕获开始和结束事件，并相应地处理链路信息span。在这些例子中，我将使用go.opencensus.io/trace提供的跟踪工具，它对接谷歌云跟踪，但其他跟踪库的行为应该类似。

现在我们有一个函数在查询开始时调用，我们需要引入链路追逐:

func beforeQuery(scope *gorm.DB) {
    db.Statement.Context = startTrace(
    db.Statement.Context,
    db,
    operation,
  )
}

func startTrace(
  ctx context.Context,
  db *gorm.DB,
) context.Context {
    // 判断是否需要启动链路追逐，查看追踪的span是否存在
    if span := trace.FromContext(ctx); span == nil {
        return ctx
    }

    ctx, span := trace.StartSpan(ctx, "gorm.query")
    return ctx
}

然后我们需要对这个追踪span收尾处理:

func afterQuery(scope *gorm.DB) { endTrace(scope) }

func endTrace(db *gorm.DB) {
    span := trace.FromContext(db.Statement.Context)
    if span == nil || !span.IsRecordingEvents() {
        return
    }

    var status trace.Status

    if db.Error != nil {
        err := db.Error
        if err == gorm.ErrRecordNotFound {
            status.Code = trace.StatusCodeNotFound
        } else {
            status.Code = trace.StatusCodeUnknown
        }

        status.Message = err.Error()
    }
    span.SetStatus(status)
    span.End()
}

db.Callback().
    Query().
    After("gorm:query").
    Register("instrumentation:after_query", afterQuery)

现在我们可以在链路追踪中看到所有gorm查询!

然而，上图中不太清楚查询实际上在做什么，让我们看看是否可以让这些span包含更多有用信息，通过添加:

• 数据库表名信息和查询指纹
• 函数调用代码行数
• 查询的WHERE参数
• 影响表行数

查询指纹是查询的唯一标识符，与格式和变量无关，因此您可以唯一地标识在数据库中具有相同行为的查询。
让我们扩展前面的代码:

func startTrace(ctx context.Context, db *gorm.DB) context.Context {
    // Don't trace queries if they don't have a parent span.
    if span := trace.FromContext(ctx); span == nil {
        return ctx
    }

    // start the span
    ctx, span := trace.StartSpan(ctx, fmt.Sprintf("gorm.query.%s", db.Statement.Table))

    // set the caller of the gorm query, so we know where in the codebase the
    // query originated.
    //
    // walk up the call stack looking for the line of code that called us. but
    // give up if it's more than 20 steps, and skip the first 5 as they're all
    // gorm anyway
    var (
        file string
        line int
    )
    for n := 5; n < 20; n++ {
        _, file, line, _ = runtime.Caller(n)
        if strings.Contains(file, "/gorm.io/") {
            // skip any helper code and go further up the call stack
            continue
        }
        break
    }
    span.AddAttributes(trace.StringAttribute("caller", fmt.Sprintf("%s:%v", file, line)))

    // add the primary table to the span metadata
    span.AddAttributes(trace.StringAttribute("gorm.table", db.Statement.Table))
    return ctx
}

func endTrace(db *gorm.DB) {
    // get the span from the context
    span := trace.FromContext(db.Statement.Context)
    if span == nil || !span.IsRecordingEvents() {
        return
    }

    // set the span status, so we know if the query was successful
    var status trace.Status
    if db.Error != nil {
        err := db.Error
        if err == gorm.ErrRecordNotFound {
            status.Code = trace.StatusCodeNotFound
        } else {
            status.Code = trace.StatusCodeUnknown
        }

        status.Message = err.Error()
    }
    span.SetStatus(status)

    // add the number of affected rows & query string to the span metadata
    span.AddAttributes(
        trace.Int64Attribute("gorm.rows_affected", db.Statement.RowsAffected),
        trace.StringAttribute("gorm.query", db.Statement.SQL.String()),
    )
    // Query fingerprint provided by github.com/pganalyze/pg_query_go
    fingerprint, err := pg_query.Fingerprint(db.Statement.SQL.String())
    if err != nil {
        fingerprint = "unknown"
    }

    // Rename the span with the fingerprint, as the DB handle
    // doesn't have SQL to fingerprint before being executed
    span.SetName(fmt.Sprintf("gorm.query.%s.%s", db.Statement.Table, fingerprint))

    // finally end the span
    span.End()
}

func afterQuery(scope *gorm.DB) {
    // now in afterQuery we can add query vars to the span metadata
    // we do this in afterQuery rather than the trace functions so we
    // can re-use the traces for non-select cases where we wouldn't want
    // to record the vars as they may contain sensitive data

    // first we extract the vars from the query & map them into a
  // human readable format
    fieldStrings := []string{}
    if scope.Statement != nil {
        fieldStrings = lo.Map(scope.Statement.Vars, func(v any i int) string {
            return fmt.Sprintf("($%v = %v)", i+1, v)
        })
    }
    // then add the vars to the span metadata
    span := trace.FromContext(scope.Statement.Context)
    if span != nil && span.IsRecordingEvents() {
        span.AddAttributes(
            trace.StringAttribute("gorm.query.vars", strings.Join(fieldStrings, ", ")),
        )
    }
    endTrace(scope)
}

现在，我们获得了非常简单详细的数据库查询跟踪信息，使我们更容易理解我们的应用程序在做什么!

Gorm为查询生命周期的不同部分提供回调，你可以为它们添加特定的行为，我们目前分别跟踪创建、删除、更新和查询，但如果你想更进一步，你可以查看Gorm文档!． 你可以在这里[https://gist.github.com/arussellsaw/bbedfdefee119b4600ce085b773da4b9]找到这篇文章中的所有代码。

请记住，如果不小心，您可能会追踪到一些敏感数据。因此，请确保清理您的查询变量。一个好的实践是只跟踪SELECT查询，因为它们通常是通过ID完成的，而不是任何敏感信息。