gorm v2 tracing

.gitignore

@@ -0,0 +1 @@
+.idea

README.md

@@ -1,11 +1,11 @@
 # opentracing gorm
 
-[OpenTracing](http://opentracing.io/) instrumentation for [GORM](http://gorm.io/).
+[OpenTracing](http://opentracing.io/) instrumentation for [GORM V2](https://github.com/go-gorm/gorm).
 
 ## Install
 
 ```
-go get -u github.com/smacker/opentracing-gorm
+go get -u github.com/lhypj/opentracing-gorm
 ```
 
 ## Usage
@@ -23,7 +23,7 @@ func init() {
 }
 
 func initDB() *gorm.DB {
-    db, err := gorm.Open("sqlite3", ":memory:")
+    instance, err := gorm.Open(mysql.New(mysql.Config{DSN: dsn}), &gorm.Config{})
     if err != nil {
         panic(err)
     }

go.mod

@@ -0,0 +1,10 @@
+module github.com/lhypj/opentracing-gorm
+
+go 1.14
+
+require (
+	github.com/opentracing/opentracing-go v1.1.0
+	github.com/stretchr/testify v1.6.1 // indirect
+	gorm.io/driver/mysql v0.2.4
+	gorm.io/gorm v0.2.18
+)

go.sum

@@ -0,0 +1,23 @@
+github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8=
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/go-sql-driver/mysql v1.5.0 h1:ozyZYNQW3x3HtqT1jira07DN2PArx2v7/mN66gGcHOs=
+github.com/go-sql-driver/mysql v1.5.0/go.mod h1:DCzpHaOWr8IXmIStZouvnhqoel9Qv2LBy8hT2VhHyBg=
+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.1 h1:g39TucaRWyV3dwDO++eEc6qf8TVIQ/Da48WmqjZ3i7E=
+github.com/jinzhu/now v1.1.1/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8=
+github.com/opentracing/opentracing-go v1.1.0 h1:pWlfV3Bxv7k65HYwkikxat0+s3pV4bsqf19k25Ur8rU=
+github.com/opentracing/opentracing-go v1.1.0/go.mod h1:UkNAQd3GIcIGf0SeVgPpRdFStlNbqXla1AfSYxPUl2o=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/testify v1.6.1 h1:hDPOHmpOpP40lSULcqw7IrRb/u7w6RpDC9399XyoNd0=
+github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+gorm.io/driver/mysql v0.2.4 h1:o8B77WK1YNhwDPjjHW1VHls0s/N0hJcEzf2NFjIm3Uk=
+gorm.io/driver/mysql v0.2.4/go.mod h1:TI8ClyQTGmgbIVmc5jKbmsrk8bKl6ewZXwePBfcENII=
+gorm.io/gorm v0.2.18 h1:Sij/WJhJFmP8mGbq1Rc2uTZQGExC0kyFQNLKBBUQiDU=
+gorm.io/gorm v0.2.18/go.mod h1:0HFTzE/SqkGTzK6TlDPPQbAYCluiVvhzoA1+aVyzenw=

otgorm.go

@@ -3,11 +3,10 @@ package otgorm
 import (
 	"context"
 	"fmt"
-	"strings"
-
-	"github.com/jinzhu/gorm"
 	opentracing "github.com/opentracing/opentracing-go"
 	"github.com/opentracing/opentracing-go/ext"
+	"gorm.io/gorm"
+	"strings"
 )
 
 const (
@@ -43,44 +42,45 @@ func newCallbacks() *callbacks {
 	return &callbacks{}
 }
 
-func (c *callbacks) beforeCreate(scope *gorm.Scope)   { c.before(scope) }
-func (c *callbacks) afterCreate(scope *gorm.Scope)    { c.after(scope, "INSERT") }
-func (c *callbacks) beforeQuery(scope *gorm.Scope)    { c.before(scope) }
-func (c *callbacks) afterQuery(scope *gorm.Scope)     { c.after(scope, "SELECT") }
-func (c *callbacks) beforeUpdate(scope *gorm.Scope)   { c.before(scope) }
-func (c *callbacks) afterUpdate(scope *gorm.Scope)    { c.after(scope, "UPDATE") }
-func (c *callbacks) beforeDelete(scope *gorm.Scope)   { c.before(scope) }
-func (c *callbacks) afterDelete(scope *gorm.Scope)    { c.after(scope, "DELETE") }
-func (c *callbacks) beforeRowQuery(scope *gorm.Scope) { c.before(scope) }
-func (c *callbacks) afterRowQuery(scope *gorm.Scope)  { c.after(scope, "") }
+func (c *callbacks) beforeCreate(db *gorm.DB)   { c.before(db) }
+func (c *callbacks) afterCreate(db *gorm.DB)    { c.after(db, "INSERT") }
+func (c *callbacks) beforeQuery(db *gorm.DB)    { c.before(db) }
+func (c *callbacks) afterQuery(db *gorm.DB)     { c.after(db, "SELECT") }
+func (c *callbacks) beforeUpdate(db *gorm.DB)   { c.before(db) }
+func (c *callbacks) afterUpdate(db *gorm.DB)    { c.after(db, "UPDATE") }
+func (c *callbacks) beforeDelete(db *gorm.DB)   { c.before(db) }
+func (c *callbacks) afterDelete(db *gorm.DB)    { c.after(db, "DELETE") }
+func (c *callbacks) beforeRowQuery(db *gorm.DB) { c.before(db) }
+func (c *callbacks) afterRowQuery(db *gorm.DB)  { c.after(db, "") }
 
-func (c *callbacks) before(scope *gorm.Scope) {
-	val, ok := scope.Get(parentSpanGormKey)
+func (c *callbacks) before(db *gorm.DB) {
+	val, ok := db.Get(parentSpanGormKey)
 	if !ok {
 		return
 	}
 	parentSpan := val.(opentracing.Span)
 	tr := parentSpan.Tracer()
 	sp := tr.StartSpan("sql", opentracing.ChildOf(parentSpan.Context()))
 	ext.DBType.Set(sp, "sql")
-	scope.Set(spanGormKey, sp)
+	db.Set(spanGormKey, sp)
 }
 
-func (c *callbacks) after(scope *gorm.Scope, operation string) {
-	val, ok := scope.Get(spanGormKey)
+func (c *callbacks) after(db *gorm.DB, operation string) {
+	val, ok := db.Get(spanGormKey)
 	if !ok {
 		return
 	}
+	sql := db.Statement.SQL.String()
 	sp := val.(opentracing.Span)
 	if operation == "" {
-		operation = strings.ToUpper(strings.Split(scope.SQL, " ")[0])
+		operation = strings.ToUpper(strings.Split(sql, " ")[0])
 	}
-	ext.Error.Set(sp, scope.HasError())
-	ext.DBStatement.Set(sp, scope.SQL)
-	sp.SetTag("db.table", scope.TableName())
+	ext.Error.Set(sp, db.Error != nil)
+	ext.DBStatement.Set(sp, sql)
+	sp.SetTag("db.table", db.Statement.Table)
 	sp.SetTag("db.method", operation)
-	sp.SetTag("db.err", scope.HasError())
-	sp.SetTag("db.count", scope.DB().RowsAffected)
+	sp.SetTag("db.err", db.Error != nil)
+	sp.SetTag("db.count", db.RowsAffected)
 	sp.Finish()
 }
 
@@ -103,7 +103,7 @@ func registerCallbacks(db *gorm.DB, name string, c *callbacks) {
 		db.Callback().Delete().Before(gormCallbackName).Register(beforeName, c.beforeDelete)
 		db.Callback().Delete().After(gormCallbackName).Register(afterName, c.afterDelete)
 	case "row_query":
-		db.Callback().RowQuery().Before(gormCallbackName).Register(beforeName, c.beforeRowQuery)
-		db.Callback().RowQuery().After(gormCallbackName).Register(afterName, c.afterRowQuery)
+		db.Callback().Row().Before(gormCallbackName).Register(beforeName, c.beforeRowQuery)
+		db.Callback().Row().Before(gormCallbackName).Register(afterName, c.afterRowQuery)
 	}
 }

otgorm_test.go

@@ -2,48 +2,49 @@ package otgorm_test
 
 import (
 	"context"
+	"log"
 	"testing"
 
-	"github.com/jinzhu/gorm"
-	_ "github.com/jinzhu/gorm/dialects/sqlite"
+	otgorm "github.com/lhypj/opentracing-gorm"
 	"github.com/opentracing/opentracing-go"
 	"github.com/opentracing/opentracing-go/mocktracer"
-	otgorm "github.com/smacker/opentracing-gorm"
+	"gorm.io/driver/mysql"
+	"gorm.io/gorm"
 )
 
 var tracer *mocktracer.MockTracer
-var gDB *gorm.DB
 
-func init() {
-	gDB = initDB()
-	tracer = mocktracer.New()
-	opentracing.SetGlobalTracer(tracer)
-}
-
-type Product struct {
-	gorm.Model
-	Code string
-}
-
-func initDB() *gorm.DB {
-	db, err := gorm.Open("sqlite3", ":memory:")
+func GetInstance() *gorm.DB {
+	var db *gorm.DB
+	dsn := "root:zxcvbnm123@tcp(localhost:3306)/testdb?parseTime=True&loc=Asia%2FShanghai"
+	db, err := gorm.Open(mysql.New(mysql.Config{DSN: dsn}), &gorm.Config{})
 	if err != nil {
-		panic(err)
+		log.Fatalf("open db: %s", err)
 	}
 	db.AutoMigrate(&Product{})
 	db.Create(&Product{Code: "L1212"})
 	otgorm.AddGormCallbacks(db)
 	return db
 }
 
+type Product struct {
+	gorm.Model
+	Code string
+}
+
 func Handler(ctx context.Context) {
+	db := GetInstance()
+	tracer = mocktracer.New()
+	opentracing.SetGlobalTracer(tracer)
+
+
 	span, ctx := opentracing.StartSpanFromContext(ctx, "handler")
 	defer span.Finish()
 
-	db := otgorm.SetSpanToGorm(ctx, gDB)
+	db = otgorm.SetSpanToGorm(ctx, db)
 
 	var product Product
-	db.First(&product, 1)
+	db.WithContext(ctx).Where("id = 1").First(&product)
 }
 
 func TestPool(t *testing.T) {
@@ -63,7 +64,7 @@ func TestPool(t *testing.T) {
 		"db.table":     "products",
 		"db.method":    "SELECT",
 		"db.type":      "sql",
-		"db.statement": `SELECT * FROM "products"  WHERE "products"."deleted_at" IS NULL AND (("products"."id" = 1)) ORDER BY "products"."id" ASC LIMIT 1`,
+		"db.statement": "SELECT * FROM `products` WHERE id = 1 AND `products`.`deleted_at` IS NULL ORDER BY `products`.`id` LIMIT 1",
 		"db.err":       false,
 		"db.count":     int64(1),
 	}
@@ -80,7 +81,7 @@ func TestPool(t *testing.T) {
 			continue
 		}
 		if value != expected {
-			t.Errorf("sql span tag '%s' should have value '%s' but it has '%s'", name, expected, value)
+			t.Errorf("sql span tag '%s' should have value '%s' but it has \n'%s'", name, expected, value)
 		}
 	}