Skip to content

Commit

Permalink
feat: Replacing custom database ddtrace with dd-trace-go library
Browse files Browse the repository at this point in the history
  • Loading branch information
@laynax
laynax committed Jun 13, 2024
1 parent bc49b52 commit 0ac4264
Show file tree
Hide file tree
Showing 4 changed files with 71 additions and 287 deletions.
16 changes: 16 additions & 0 deletions go.sum
View file
Original file line number Diff line number Diff line change
Expand Up @@ -149,6 +149,10 @@ github.com/gobwas/ws v1.0.2/go.mod h1:szmBTxLgaFppYjEmNtny/v3w89xOydFnnZMcgRRu/E
github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q=
github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q=
github.com/golang-jwt/jwt v3.2.2+incompatible/go.mod h1:8pz2t5EyA70fFQQSrl6XZXzqecmYZeUEB8OUGHkxJ+I=
github.com/golang-sql/civil v0.0.0-20220223132316-b832511892a9 h1:au07oEsX2xN0ktxqI+Sida1w446QrXBRJ0nee3SNZlA=
github.com/golang-sql/civil v0.0.0-20220223132316-b832511892a9/go.mod h1:8vg3r2VgvsThLBIFL93Qb5yWzgyZWhEmBwUJWevAkK0=
github.com/golang-sql/sqlexp v0.1.0 h1:ZCD6MBpcuOVfGVqsEmY5/4FtYiKz6tSyUv9LPEDei6A=
github.com/golang-sql/sqlexp v0.1.0/go.mod h1:J4ad9Vo8ZCWQ2GMrC4UCQy1JpCbwU9m3EOqtpKwwwHI=
github.com/golang/mock v1.6.0/go.mod h1:p6yTPP+5HYm5mzsMV8JkE6ZKdX+/wYM6Hr+LicevLPs=
github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
github.com/golang/protobuf v1.3.1/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
Expand Down Expand Up @@ -193,6 +197,12 @@ github.com/iris-contrib/go.uuid v2.0.0+incompatible/go.mod h1:iz2lgM/1UnEf1kP0L/
github.com/iris-contrib/jade v1.1.3/go.mod h1:H/geBymxJhShH5kecoiOCSssPX7QWYH7UaeZTSWddIk=
github.com/iris-contrib/pongo2 v0.0.1/go.mod h1:Ssh+00+3GAZqSQb30AvBRNxBx7rf0GqwkjqxNd0u65g=
github.com/iris-contrib/schema v0.0.1/go.mod h1:urYA3uvUNG1TIIjOSCzHr9/LmbQo8LrOcOqfqxa4hXw=
github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg=
github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a h1:bbPeKD0xmW/Y25WS6cokEszi5g+S0QxI/d45PkRi7Nk=
github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM=
github.com/jackc/pgx/v5 v5.4.2 h1:u1gmGDwbdRUZiwisBm/Ky2M14uQyUP65bG8+20nnyrg=
github.com/jackc/pgx/v5 v5.4.2/go.mod h1:q6iHT8uDNXWiFNOlRqJzBTaSH3+2xCXkokxHZC5qWFY=
github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E=
github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc=
github.com/jinzhu/now v1.1.4/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8=
Expand Down Expand Up @@ -248,6 +258,8 @@ github.com/mattn/go-sqlite3 v1.14.22/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxU
github.com/mattn/goveralls v0.0.2/go.mod h1:8d1ZMHsd7fW6IRPKQh46F2WRpyib5/X4FOpevwGNQEw=
github.com/mediocregopher/radix/v3 v3.4.2/go.mod h1:8FL3F6UQRXHXIBSPUs5h0RybMF8i4n7wVopoX3x7Bv8=
github.com/microcosm-cc/bluemonday v1.0.2/go.mod h1:iVP4YcDBq+n/5fb23BhYFvIMq/leAFZyRl6bYmGDlGc=
github.com/microsoft/go-mssqldb v0.21.0 h1:p2rpHIL7TlSv1QrbXJUAcbyRKnIT0C9rRkH2E4OjLn8=
github.com/microsoft/go-mssqldb v0.21.0/go.mod h1:+4wZTUnz/SV6nffv+RRRB/ss8jPng5Sho2SmM1l2ts4=
github.com/mitchellh/go-homedir v1.1.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0=
github.com/mitchellh/mapstructure v1.1.2/go.mod h1:FVVH3fgwuzCH5S8UJGiWEs2h04kUh9fWfEaFds41c1Y=
github.com/mitchellh/mapstructure v1.5.0 h1:jeMsZIYE/09sWLaz43PL7Gy6RuMjD2eJVyuac5Z2hdY=
Expand Down Expand Up @@ -508,8 +520,12 @@ gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
gorm.io/driver/mysql v1.4.6 h1:5zS3vIKcyb46byXZNcYxaT9EWNIhXzu0gPuvvVrwZ8s=
gorm.io/driver/mysql v1.4.6/go.mod h1:SxzItlnT1cb6e1e4ZRpgJN2VYtcqJgqnHxWr4wsP8oc=
gorm.io/driver/postgres v1.4.6 h1:1FPESNXqIKG5JmraaH2bfCVlMQ7paLoCreFxDtqzwdc=
gorm.io/driver/postgres v1.4.6/go.mod h1:UJChCNLFKeBqQRE+HrkFUbKbq9idPXmTOk2u4Wok8S4=
gorm.io/driver/sqlite v1.4.4 h1:gIufGoR0dQzjkyqDyYSCvsYR6fba1Gw5YKDqKeChxFc=
gorm.io/driver/sqlite v1.4.4/go.mod h1:0Aq3iPO+v9ZKbcdiz8gLWRw5VOPcBOPUQJFLq5e2ecI=
gorm.io/driver/sqlserver v1.4.2 h1:nMtEeKqv2R/vv9FoHUFWfXfP6SskAgRar0TPlZV1stk=
gorm.io/driver/sqlserver v1.4.2/go.mod h1:XHwBuB4Tlh7DqO0x7Ema8dmyWsQW7wi38VQOAFkrbXY=
gorm.io/gorm v1.23.8/go.mod h1:l2lP/RyAtc1ynaTjFksBde/O8v9oOGIApu2/xRitmZk=
gorm.io/gorm v1.24.0/go.mod h1:DVrVomtaYTbqs7gB/x2uVvqnXzv0nqjB396B8cG4dBA=
gorm.io/gorm v1.25.3 h1:zi4rHZj1anhZS2EuEODMhDisGy+Daq9jtPrNGgbQYD8=
Expand Down
45 changes: 30 additions & 15 deletions pkg/database/gorm.go
View file
Original file line number Diff line number Diff line change
@@ -1,18 +1,36 @@
package database

import (
"fmt"
"strconv"
"time"

"github.com/DATA-DOG/go-txdb"
gormtrace "gopkg.in/DataDog/dd-trace-go.v1/contrib/gorm.io/gorm.v1"
"gorm.io/driver/mysql"
"gorm.io/gorm"
)

const testEnv = "test"

// NewConnection returns a new Gorm database connection.
func NewConnection(config *Config, environment string) (*gorm.DB, error) {
// NewConnection returns a new instrumented Gorm database connection.
func NewConnection(config *Config, environment, appName string) (*gorm.DB, error) {
serviceName := fmt.Sprintf("%s-mysql", appName)
dialector := getDialectorFromConfig(config, environment)

db, err := gormtrace.Open(dialector, nil, gormtrace.WithServiceName(serviceName))
if err != nil {
return nil, err
}

if err := databasePoolSettings(db, config); err != nil {
return nil, err
}

return db, nil
}

func getDialectorFromConfig(config *Config, environment string) gorm.Dialector {
connectionDetails := NewConnectionDetails(config)

connectionString := connectionDetails.String()
Expand All @@ -29,25 +47,22 @@ func NewConnection(config *Config, environment string) (*gorm.DB, error) {
connectionString = testDriverName
}

dialector := mysql.New(mysql.Config{
return mysql.New(mysql.Config{
DSN: connectionString,
DriverName: driverName,
})
}

db, err := gorm.Open(dialector)
if err != nil {
return nil, err
}

sqlDB, err := db.DB()
func databasePoolSettings(gormDB *gorm.DB, config *Config) error {
db, err := gormDB.DB()
if err != nil {
return nil, err
return err
}

sqlDB.SetMaxIdleConns(config.Pool)
sqlDB.SetMaxOpenConns(config.MaxOpenConnections)
sqlDB.SetConnMaxIdleTime(config.ConnectionMaxIdleTime)
sqlDB.SetConnMaxLifetime(config.ConnectionMaxLifetime)
db.SetMaxIdleConns(config.Pool)
db.SetMaxOpenConns(config.MaxOpenConnections)
db.SetConnMaxIdleTime(config.ConnectionMaxIdleTime)
db.SetConnMaxLifetime(config.ConnectionMaxLifetime)

return db, nil
return nil
}
151 changes: 5 additions & 146 deletions pkg/instrumentation/database.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,159 +2,18 @@ package instrumentation

import (
"context"
"fmt"
"strings"

"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
ddtrace "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
"gorm.io/gorm"
)

type spanContextKey string

var (
// ParentSpanGormKey is the name of the parent span key
parentSpanGormKey = spanContextKey("trancingParentSpan")
// SpanGormKey is the name of the span key
spanGormKey = spanContextKey("tracingSpan")
)

// TraceDatabase sets span to gorm settings, returns cloned DB
// TraceDatabase fetches the span from the context and injects it to a new
// database session's statement context.
func TraceDatabase(ctx context.Context, db *gorm.DB) *gorm.DB {
if ctx == nil {
return db
}

parentSpan, _ := ddtrace.SpanFromContext(ctx)

return db.Session(&gorm.Session{
Context: context.WithValue(db.Statement.Context, parentSpanGormKey, parentSpan),
})
}

// InstrumentDatabase adds callbacks for tracing, call TraceDatabase to make it work
func InstrumentDatabase(db *gorm.DB, appName string) {
callbacks := newCallbacks(appName)

registerCallbacks(db, "create", callbacks)
registerCallbacks(db, "query", callbacks)
registerCallbacks(db, "update", callbacks)
registerCallbacks(db, "delete", callbacks)
registerCallbacks(db, "row", callbacks)
}

type callbacks struct {
serviceName string
}

func newCallbacks(appName string) *callbacks {
return &callbacks{
serviceName: fmt.Sprintf("%s-%s", appName, "mysql"),
}
}

func (c *callbacks) beforeCreate(db *gorm.DB) { c.before(db, "INSERT", c.serviceName) }
func (c *callbacks) afterCreate(db *gorm.DB) { c.after(db) }
func (c *callbacks) beforeQuery(db *gorm.DB) { c.before(db, "SELECT", c.serviceName) }
func (c *callbacks) afterQuery(db *gorm.DB) { c.after(db) }
func (c *callbacks) beforeUpdate(db *gorm.DB) { c.before(db, "UPDATE", c.serviceName) }
func (c *callbacks) afterUpdate(db *gorm.DB) { c.after(db) }
func (c *callbacks) beforeDelete(db *gorm.DB) { c.before(db, "DELETE", c.serviceName) }
func (c *callbacks) afterDelete(db *gorm.DB) { c.after(db) }
func (c *callbacks) beforeRow(db *gorm.DB) { c.before(db, "", c.serviceName) }
func (c *callbacks) afterRow(db *gorm.DB) { c.after(db) }
func (c *callbacks) before(db *gorm.DB, operationName string, serviceName string) {
if db.Statement == nil || db.Statement.Context == nil {
return
}

parentSpan, ok := db.Statement.Context.Value(parentSpanGormKey).(ddtrace.Span)
if !ok {
return
}

spanOpts := []ddtrace.StartSpanOption{
ddtrace.ChildOf(parentSpan.Context()),
ddtrace.SpanType(ext.SpanTypeSQL),
ddtrace.ServiceName(serviceName),
}
if operationName == "" {
operationName = strings.Split(db.Statement.SQL.String(), " ")[0]
}
sp := ddtrace.StartSpan(operationName, spanOpts...)
db.Statement.Context = context.WithValue(db.Statement.Context, spanGormKey, sp)
}

func (c *callbacks) after(db *gorm.DB) {
if db.Statement == nil || db.Statement.Context == nil {
return
}

sp, ok := db.Statement.Context.Value(spanGormKey).(ddtrace.Span)
if !ok {
return
}

sp.SetTag(ext.ResourceName, strings.ToUpper(db.Statement.SQL.String()))
sp.SetTag("db.table", db.Statement.Table)
sp.SetTag("db.query", strings.ToUpper(db.Statement.SQL.String()))
sp.SetTag("db.err", db.Error)
sp.SetTag("db.count", db.RowsAffected)
sp.Finish()
}

func registerCallbacks(db *gorm.DB, name string, c *callbacks) {
var err error

beforeName := fmt.Sprintf("tracing:%v_before", name)
afterName := fmt.Sprintf("tracing:%v_after", name)
gormCallbackName := fmt.Sprintf("gorm:%v", name)
// gorm does some magic, if you pass CallbackProcessor here - nothing works
switch name {
case "create":
err = db.Callback().Create().Before(gormCallbackName).Register(beforeName, c.beforeCreate)
if err != nil {
return
}
err = db.Callback().Create().After(gormCallbackName).Register(afterName, c.afterCreate)
if err != nil {
return
}
case "query":
err = db.Callback().Query().Before(gormCallbackName).Register(beforeName, c.beforeQuery)
if err != nil {
return
}
err = db.Callback().Query().After(gormCallbackName).Register(afterName, c.afterQuery)
if err != nil {
return
}
case "update":
err = db.Callback().Update().Before(gormCallbackName).Register(beforeName, c.beforeUpdate)
if err != nil {
return
}
err = db.Callback().Update().After(gormCallbackName).Register(afterName, c.afterUpdate)
if err != nil {
return
}
case "delete":
err = db.Callback().Delete().Before(gormCallbackName).Register(beforeName, c.beforeDelete)
if err != nil {
return
}
err = db.Callback().Delete().After(gormCallbackName).Register(afterName, c.afterDelete)
if err != nil {
return
}
case "row":
err = db.Callback().Row().Before(gormCallbackName).Register(beforeName, c.beforeRow)
if err != nil {
return
}
err = db.Callback().Row().After(gormCallbackName).Register(afterName, c.afterRow)
if err != nil {
return
}
}
parentSpan, _ := tracer.SpanFromContext(ctx)
return db.WithContext(tracer.ContextWithSpan(ctx, parentSpan))
}
Loading

0 comments on commit 0ac4264

Please sign in to comment.