Skip to content

Commit

Permalink
Merge branch 'main' into dario.castane/bun-contrib
Browse files Browse the repository at this point in the history
  • Loading branch information
@darccio
darccio authored Jul 8, 2024
2 parents 33f5b50 + 2664189 commit 6c08e35
Show file tree
Hide file tree
Showing 4 changed files with 217 additions and 45 deletions.
26 changes: 26 additions & 0 deletions contrib/gorm.io/gorm.v1/example_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ package gorm_test

import (
"context"
"errors"
"log"

sqltrace "gopkg.in/DataDog/dd-trace-go.v1/contrib/database/sql"
Expand Down Expand Up @@ -41,6 +42,31 @@ func ExampleOpen() {
db.Where("name = ?", "jinzhu").First(&user)
}

// ExampleNewTracePlugin illustrates how to trace gorm using the gorm.Plugin api.
func ExampleNewTracePlugin() {
// Register augments the provided driver with tracing, enabling it to be loaded by gorm.Open and the gormtrace.TracePlugin.
sqltrace.Register("pgx", &stdlib.Driver{}, sqltrace.WithServiceName("my-service"))
sqlDb, err := sqltrace.Open("pgx", "postgres://pqgotest:password@localhost/pqgotest?sslmode=disable")
if err != nil {
log.Fatal(err)
}
db, err := gorm.Open(postgres.New(postgres.Config{Conn: sqlDb}), &gorm.Config{})
if err != nil {
log.Fatal(err)
}
var user User

errCheck := gormtrace.WithErrorCheck(func(err error) bool {
return !errors.Is(err, gorm.ErrRecordNotFound)
})
if err := db.Use(gormtrace.NewTracePlugin(errCheck)); err != nil {
log.Fatal(err)
}

// All calls through gorm.DB are now traced.
db.Where("name = ?", "jinzhu").First(&user)
}

func ExampleContext() {
// Register augments the provided driver with tracing, enabling it to be loaded by gormtrace.Open.
sqltrace.Register("pgx", &stdlib.Driver{}, sqltrace.WithServiceName("my-service"))
Expand Down
96 changes: 61 additions & 35 deletions contrib/gorm.io/gorm.v1/gorm.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,7 @@
package gorm

import (
"context"
"math"
"time"

"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
Expand All @@ -33,6 +31,26 @@ const (
gormSpanStartTimeKey = key("dd-trace-go:span")
)

type tracePlugin struct {
options []Option
}

// NewTracePlugin returns a new gorm.Plugin that enhances the underlying *gorm.DB with tracing.
func NewTracePlugin(opts ...Option) gorm.Plugin {
return tracePlugin{
options: opts,
}
}

func (tracePlugin) Name() string {
return "DDTracePlugin"
}

func (g tracePlugin) Initialize(db *gorm.DB) error {
_, err := withCallbacks(db, g.options...)
return err
}

// Open opens a new (traced) database connection. The used driver must be formerly registered
// using (gopkg.in/DataDog/dd-trace-go.v1/contrib/database/sql).Register.
func Open(dialector gorm.Dialector, cfg *gorm.Config, opts ...Option) (*gorm.DB, error) {
Expand All @@ -57,86 +75,80 @@ func withCallbacks(db *gorm.DB, opts ...Option) (*gorm.DB, error) {
}
log.Debug("Registering Callbacks: %#v", cfg)

afterFunc := func(operationName string) func(*gorm.DB) {
afterFunc := func() func(*gorm.DB) {
return func(db *gorm.DB) {
after(db, cfg)
}
}

beforeFunc := func(operationName string) func(*gorm.DB) {
return func(db *gorm.DB) {
after(db, operationName, cfg)
before(db, operationName, cfg)
}
}

cb := db.Callback()
err := cb.Create().Before("gorm:create").Register("dd-trace-go:before_create", before)
err := cb.Create().Before("gorm:create").Register("dd-trace-go:before_create", beforeFunc("gorm.create"))
if err != nil {
return db, err
}
err = cb.Create().After("gorm:create").Register("dd-trace-go:after_create", afterFunc("gorm.create"))
err = cb.Create().After("gorm:create").Register("dd-trace-go:after_create", afterFunc())
if err != nil {
return db, err
}
err = cb.Update().Before("gorm:update").Register("dd-trace-go:before_update", before)
err = cb.Update().Before("gorm:update").Register("dd-trace-go:before_update", beforeFunc("gorm.update"))
if err != nil {
return db, err
}
err = cb.Update().After("gorm:update").Register("dd-trace-go:after_update", afterFunc("gorm.update"))
err = cb.Update().After("gorm:update").Register("dd-trace-go:after_update", afterFunc())
if err != nil {
return db, err
}
err = cb.Delete().Before("gorm:delete").Register("dd-trace-go:before_delete", before)
err = cb.Delete().Before("gorm:delete").Register("dd-trace-go:before_delete", beforeFunc("gorm.delete"))
if err != nil {
return db, err
}
err = cb.Delete().After("gorm:delete").Register("dd-trace-go:after_delete", afterFunc("gorm.delete"))
err = cb.Delete().After("gorm:delete").Register("dd-trace-go:after_delete", afterFunc())
if err != nil {
return db, err
}
err = cb.Query().Before("gorm:query").Register("dd-trace-go:before_query", before)
err = cb.Query().Before("gorm:query").Register("dd-trace-go:before_query", beforeFunc("gorm.query"))
if err != nil {
return db, err
}
err = cb.Query().After("gorm:query").Register("dd-trace-go:after_query", afterFunc("gorm.query"))
err = cb.Query().After("gorm:query").Register("dd-trace-go:after_query", afterFunc())
if err != nil {
return db, err
}
err = cb.Row().Before("gorm:query").Register("dd-trace-go:before_row_query", before)
err = cb.Row().Before("gorm:row").Register("dd-trace-go:before_row_query", beforeFunc("gorm.row_query"))
if err != nil {
return db, err
}
err = cb.Row().After("gorm:query").Register("dd-trace-go:after_row_query", afterFunc("gorm.row_query"))
err = cb.Row().After("gorm:row").Register("dd-trace-go:after_row_query", afterFunc())
if err != nil {
return db, err
}
err = cb.Raw().Before("gorm:query").Register("dd-trace-go:before_raw_query", before)
err = cb.Raw().Before("gorm:raw").Register("dd-trace-go:before_raw_query", beforeFunc("gorm.raw_query"))
if err != nil {
return db, err
}
err = cb.Raw().After("gorm:query").Register("dd-trace-go:after_raw_query", afterFunc("gorm.raw_query"))
err = cb.Raw().After("gorm:raw").Register("dd-trace-go:after_raw_query", afterFunc())
if err != nil {
return db, err
}
return db, nil
}

func before(scope *gorm.DB) {
if scope.Statement != nil && scope.Statement.Context != nil {
scope.Statement.Context = context.WithValue(scope.Statement.Context, gormSpanStartTimeKey, time.Now())
}
}

func after(db *gorm.DB, operationName string, cfg *config) {
func before(db *gorm.DB, operationName string, cfg *config) {
if db.Statement == nil || db.Statement.Context == nil {
return
}

ctx := db.Statement.Context
t, ok := ctx.Value(gormSpanStartTimeKey).(time.Time)
if !ok {
if db.Config == nil || db.Config.DryRun {
return
}

opts := []ddtrace.StartSpanOption{
tracer.StartTime(t),
tracer.ServiceName(cfg.serviceName),
tracer.SpanType(ext.SpanTypeSQL),
tracer.ResourceName(db.Statement.SQL.String()),
tracer.Tag(ext.Component, componentName),
}
if !math.IsNaN(cfg.analyticsRate) {
Expand All @@ -148,10 +160,24 @@ func after(db *gorm.DB, operationName string, cfg *config) {
}
}

span, _ := tracer.StartSpanFromContext(ctx, operationName, opts...)
var dbErr error
if cfg.errCheck(db.Error) {
dbErr = db.Error
_, ctx := tracer.StartSpanFromContext(db.Statement.Context, operationName, opts...)
db.Statement.Context = ctx
}

func after(db *gorm.DB, cfg *config) {
if db.Statement == nil || db.Statement.Context == nil {
return
}
if db.Config == nil || db.Config.DryRun {
return
}
span, ok := tracer.SpanFromContext(db.Statement.Context)
if ok {
var dbErr error
if cfg.errCheck(db.Error) {
dbErr = db.Error
}
span.SetTag(ext.ResourceName, db.Statement.SQL.String())
span.Finish(tracer.WithError(dbErr))
}
span.Finish(tracer.WithError(dbErr))
}
Loading

0 comments on commit 6c08e35

Please sign in to comment.