Skip to content

Commit

Permalink
feat: add nrgorm apm integration
Browse files Browse the repository at this point in the history
  • Loading branch information
XiXiangFiles committed Jan 24, 2025
1 parent 5867ad9 commit e1a1229
Show file tree
Hide file tree
Showing 3 changed files with 130 additions and 0 deletions.
25 changes: 25 additions & 0 deletions v3/integrations/nrgorm/example/main.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
package main

import (
"log"
"time"

"github.com/newrelic/go-agent/v3/integrations/nrgorm"
"gorm.io/driver/postgres"
"gorm.io/gorm"
)

func main() {
dsn := "host=localhost user=gorm password=gorm dbname=gorm port=9920 sslmode=disable TimeZone=Asia/Shanghai"
_, err := gorm.Open(postgres.Open(dsn), &gorm.Config{
NowFunc: func() time.Time {
return time.Now().UTC()
},
Plugins: map[string]gorm.Plugin{
nrgorm.APMPlugin{}.Name(): nrgorm.APMPlugin{},
},
})
if err != nil {
log.Printf("gorm open failed: %v", err)
}
}
11 changes: 11 additions & 0 deletions v3/integrations/nrgorm/go.mod
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
module github.com/newrelic/go-agent/v3/integrations/nrgorm

go 1.21

require (
gorm.io/gorm v1.25.12
github.com/newrelic/go-agent/v3 v3.36.0
)


replace github.com/newrelic/go-agent/v3 => ../..
94 changes: 94 additions & 0 deletions v3/integrations/nrgorm/nrgorm.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,94 @@
package nrgorm

import (
"fmt"
"reflect"
"strconv"

"github.com/newrelic/go-agent/v3/newrelic"
"gorm.io/gorm"
)

const startTimeKey = "newrelic_start_time"

type APMPlugin struct{}

var _ gorm.Plugin = (*APMPlugin)(nil)

func (p APMPlugin) Name() string {
return "newrelic"
}

func (p APMPlugin) Initialize(db *gorm.DB) error {
err := db.Callback().Create().Before("gorm:create").Register("newrelic:before_create", beforeCallback)
if err != nil {
return err
}
err = db.Callback().Query().Before("gorm:query").Register("newrelic:before_query", beforeCallback)
if err != nil {
return err
}
err = db.Callback().Update().Before("gorm:update").Register("newrelic:before_update", beforeCallback)
if err != nil {
return err
}
err = db.Callback().Delete().Before("gorm:delete").Register("newrelic:before_delete", beforeCallback)
if err != nil {
return err
}

err = db.Callback().Create().After("gorm:create").Register("newrelic:after_create", afterCallback("INSERT"))
if err != nil {
return err
}
err = db.Callback().Query().After("gorm:query").Register("newrelic:after_query", afterCallback("SELECT"))
if err != nil {
return err
}
err = db.Callback().Update().After("gorm:update").Register("newrelic:after_update", afterCallback("UPDATE"))
if err != nil {
return err
}
err = db.Callback().Delete().After("gorm:delete").Register("newrelic:after_delete", afterCallback("DELETE"))
if err != nil {
return err
}

return nil
}

func beforeCallback(db *gorm.DB) {
if txn := newrelic.FromContext(db.Statement.Context); txn != nil {
db.Set(startTimeKey, txn.StartSegmentNow())
}
}

func afterCallback(operation string) func(db *gorm.DB) {
return func(db *gorm.DB) {
if startTime, ok := db.Get(startTimeKey); ok {
segment := newrelic.DatastoreSegment{
Product: newrelic.DatastorePostgres,
Collection: db.Statement.Table,
Operation: operation,
StartTime: startTime.(newrelic.SegmentStartTime),
ParameterizedQuery: db.Statement.SQL.String(),
QueryParameters: parseVars(db.Statement.Vars),
}
segment.End()
}
}
}

func parseVars(vars []interface{}) map[string]interface{} {
queryParameters := make(map[string]interface{})
for i, v := range vars {
i := i
v := v
val := reflect.ValueOf(v)
if val.Kind() == reflect.Ptr {
val = val.Elem()
}
queryParameters[fmt.Sprintf("$%v", strconv.Itoa(i+1))] = fmt.Sprintf("%v", val.Interface())
}
return queryParameters
}

0 comments on commit e1a1229

Please sign in to comment.