mirror of
https://github.com/superseriousbusiness/gotosocial.git
synced 2024-12-05 01:52:46 +00:00
189 lines
4.2 KiB
Go
189 lines
4.2 KiB
Go
package bunotel
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"runtime"
|
|
"strings"
|
|
"time"
|
|
|
|
"go.opentelemetry.io/otel"
|
|
"go.opentelemetry.io/otel/attribute"
|
|
"go.opentelemetry.io/otel/codes"
|
|
"go.opentelemetry.io/otel/metric/global"
|
|
"go.opentelemetry.io/otel/metric/instrument"
|
|
semconv "go.opentelemetry.io/otel/semconv/v1.12.0"
|
|
"go.opentelemetry.io/otel/trace"
|
|
|
|
"github.com/uptrace/bun"
|
|
"github.com/uptrace/bun/dialect"
|
|
"github.com/uptrace/bun/schema"
|
|
"github.com/uptrace/opentelemetry-go-extra/otelsql"
|
|
)
|
|
|
|
var (
|
|
tracer = otel.Tracer("github.com/uptrace/bun")
|
|
meter = global.Meter("github.com/uptrace/bun")
|
|
|
|
queryHistogram, _ = meter.Int64Histogram(
|
|
"go.sql.query_timing",
|
|
instrument.WithDescription("Timing of processed queries"),
|
|
instrument.WithUnit("milliseconds"),
|
|
)
|
|
)
|
|
|
|
type QueryHook struct {
|
|
attrs []attribute.KeyValue
|
|
formatQueries bool
|
|
}
|
|
|
|
var _ bun.QueryHook = (*QueryHook)(nil)
|
|
|
|
func NewQueryHook(opts ...Option) *QueryHook {
|
|
h := new(QueryHook)
|
|
for _, opt := range opts {
|
|
opt(h)
|
|
}
|
|
return h
|
|
}
|
|
|
|
func (h *QueryHook) Init(db *bun.DB) {
|
|
labels := make([]attribute.KeyValue, 0, len(h.attrs)+1)
|
|
labels = append(labels, h.attrs...)
|
|
if sys := dbSystem(db); sys.Valid() {
|
|
labels = append(labels, sys)
|
|
}
|
|
|
|
otelsql.ReportDBStatsMetrics(db.DB, otelsql.WithAttributes(labels...))
|
|
}
|
|
|
|
func (h *QueryHook) BeforeQuery(ctx context.Context, event *bun.QueryEvent) context.Context {
|
|
ctx, _ = tracer.Start(ctx, "", trace.WithSpanKind(trace.SpanKindClient))
|
|
return ctx
|
|
}
|
|
|
|
func (h *QueryHook) AfterQuery(ctx context.Context, event *bun.QueryEvent) {
|
|
operation := event.Operation()
|
|
dbOperation := semconv.DBOperationKey.String(operation)
|
|
|
|
labels := make([]attribute.KeyValue, 0, len(h.attrs)+2)
|
|
labels = append(labels, h.attrs...)
|
|
labels = append(labels, dbOperation)
|
|
if event.IQuery != nil {
|
|
if tableName := event.IQuery.GetTableName(); tableName != "" {
|
|
labels = append(labels, semconv.DBSQLTableKey.String(tableName))
|
|
}
|
|
}
|
|
|
|
queryHistogram.Record(ctx, time.Since(event.StartTime).Milliseconds(), labels...)
|
|
|
|
span := trace.SpanFromContext(ctx)
|
|
if !span.IsRecording() {
|
|
return
|
|
}
|
|
|
|
span.SetName(operation)
|
|
defer span.End()
|
|
|
|
query := h.eventQuery(event)
|
|
fn, file, line := funcFileLine("github.com/uptrace/bun")
|
|
|
|
attrs := make([]attribute.KeyValue, 0, 10)
|
|
attrs = append(attrs, h.attrs...)
|
|
attrs = append(attrs,
|
|
dbOperation,
|
|
semconv.DBStatementKey.String(query),
|
|
semconv.CodeFunctionKey.String(fn),
|
|
semconv.CodeFilepathKey.String(file),
|
|
semconv.CodeLineNumberKey.Int(line),
|
|
)
|
|
|
|
if sys := dbSystem(event.DB); sys.Valid() {
|
|
attrs = append(attrs, sys)
|
|
}
|
|
if event.Result != nil {
|
|
if n, _ := event.Result.RowsAffected(); n > 0 {
|
|
attrs = append(attrs, attribute.Int64("db.rows_affected", n))
|
|
}
|
|
}
|
|
|
|
switch event.Err {
|
|
case nil, sql.ErrNoRows, sql.ErrTxDone:
|
|
// ignore
|
|
default:
|
|
span.RecordError(event.Err)
|
|
span.SetStatus(codes.Error, event.Err.Error())
|
|
}
|
|
|
|
span.SetAttributes(attrs...)
|
|
}
|
|
|
|
func funcFileLine(pkg string) (string, string, int) {
|
|
const depth = 16
|
|
var pcs [depth]uintptr
|
|
n := runtime.Callers(3, pcs[:])
|
|
ff := runtime.CallersFrames(pcs[:n])
|
|
|
|
var fn, file string
|
|
var line int
|
|
for {
|
|
f, ok := ff.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
fn, file, line = f.Function, f.File, f.Line
|
|
if !strings.Contains(fn, pkg) {
|
|
break
|
|
}
|
|
}
|
|
|
|
if ind := strings.LastIndexByte(fn, '/'); ind != -1 {
|
|
fn = fn[ind+1:]
|
|
}
|
|
|
|
return fn, file, line
|
|
}
|
|
|
|
func (h *QueryHook) eventQuery(event *bun.QueryEvent) string {
|
|
const softQueryLimit = 8000
|
|
const hardQueryLimit = 16000
|
|
|
|
var query string
|
|
|
|
if h.formatQueries && len(event.Query) <= softQueryLimit {
|
|
query = event.Query
|
|
} else {
|
|
query = unformattedQuery(event)
|
|
}
|
|
|
|
if len(query) > hardQueryLimit {
|
|
query = query[:hardQueryLimit]
|
|
}
|
|
|
|
return query
|
|
}
|
|
|
|
func unformattedQuery(event *bun.QueryEvent) string {
|
|
if event.IQuery != nil {
|
|
if b, err := event.IQuery.AppendQuery(schema.NewNopFormatter(), nil); err == nil {
|
|
return bytesToString(b)
|
|
}
|
|
}
|
|
return string(event.QueryTemplate)
|
|
}
|
|
|
|
func dbSystem(db *bun.DB) attribute.KeyValue {
|
|
switch db.Dialect().Name() {
|
|
case dialect.PG:
|
|
return semconv.DBSystemPostgreSQL
|
|
case dialect.MySQL:
|
|
return semconv.DBSystemMySQL
|
|
case dialect.SQLite:
|
|
return semconv.DBSystemSqlite
|
|
case dialect.MSSQL:
|
|
return semconv.DBSystemMSSQL
|
|
default:
|
|
return attribute.KeyValue{}
|
|
}
|
|
}
|