// Copyright 2020 getensh.com. All rights reserved. // Use of this source code is governed by getensh.com. /** * @Author: mac * @Description: gorm opentracing plugins * @Date: 2020/3/25 10:36 */ package trace import ( "context" //"fmt" //"strings" //"github.com/opentracing/opentracing-go/log" opentracing "github.com/opentracing/opentracing-go" "gorm.io/gorm" //"github.com/opentracing/opentracing-go/ext" ) const ( parentSpanGormKey = "opentracingParentSpan" spanGormKey = "opentracingSpan" ) // SetSpanToGorm sets span to gorm settings, returns cloned DB func SetSpanToGorm(ctx context.Context, db *gorm.DB) *gorm.DB { if ctx == nil { return db } parentSpan := opentracing.SpanFromContext(ctx) if parentSpan == nil { return db } return db.Set(parentSpanGormKey, parentSpan) } // AddGormCallbacks adds callbacks for tracing, you should call SetSpanToGorm to make them work func AddGormCallbacks(db *gorm.DB) { //callbacks := newCallbacks() //registerCallbacks(db, "create", callbacks) //registerCallbacks(db, "query", callbacks) //registerCallbacks(db, "update", callbacks) //registerCallbacks(db, "delete", callbacks) //registerCallbacks(db, "row_query", callbacks) } type callbacks struct{} func newCallbacks() *callbacks { return &callbacks{} } // //func (c *callbacks) beforeCreate(scope *gorm.Scope) { c.before(scope, "INSERT") } //func (c *callbacks) afterCreate(scope *gorm.Scope) { c.after(scope, "INSERT") } //func (c *callbacks) beforeQuery(scope *gorm.Scope) { c.before(scope, "SELECT") } //func (c *callbacks) afterQuery(scope *gorm.Scope) { c.after(scope, "SELECT") } //func (c *callbacks) beforeUpdate(scope *gorm.Scope) { c.before(scope, "UPDATE") } //func (c *callbacks) afterUpdate(scope *gorm.Scope) { c.after(scope, "UPDATE") } //func (c *callbacks) beforeDelete(scope *gorm.Scope) { c.before(scope, "DELETE") } //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) before(scope *gorm.Scope, operation string) { // val, ok := scope.Get(parentSpanGormKey) // if !ok { // return // } // parentSpan := val.(opentracing.Span) // tr := parentSpan.Tracer() // // 如果是其他sql类型,将下面两个mysql改为对应的类型即可 // if operation == "" { // operation = strings.ToUpper(strings.Split(scope.SQL, " ")[0]) // } // sp := tr.StartSpan("MYSQL:"+operation, opentracing.ChildOf(parentSpan.Context())) // ext.DBType.Set(sp, "mysql") // scope.Set(spanGormKey, sp) //} // //func (c *callbacks) after(scope *gorm.DB, operation string) { // val, ok := scope.Get(spanGormKey) // if !ok { // return // } // sp := val.(opentracing.Span) // if operation == "" { // operation = strings.ToUpper(strings.Split(scope.SQL, " ")[0]) // } // if scope.HasError() { // ext.Error.Set(sp, scope.HasError()) // //ext.DBStatement.Set(sp, scope.SQL) // sp.LogFields( // log.String("db.table", scope.TableName()), // log.String("db.method", operation), // log.String("db.sql", scope.SQL), // log.Object("db.err", scope.DB().Error), // log.Int64("db.count", scope.DB().RowsAffected), // ) // } else { // sp.LogFields(log.String("db.sql", scope.SQL), // log.Int64("db.count", scope.DB().RowsAffected)) // } // // sp.Finish() //} // //func registerCallbacks(db *gorm.DB, name string, c *callbacks) { // 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": // db.Callback().Create().Before(gormCallbackName).Register(beforeName, c.beforeCreate) // db.Callback().Create().After(gormCallbackName).Register(afterName, c.afterCreate) // case "query": // db.Callback().Query().Before(gormCallbackName).Register(beforeName, c.beforeQuery) // db.Callback().Query().After(gormCallbackName).Register(afterName, c.afterQuery) // case "update": // db.Callback().Update().Before(gormCallbackName).Register(beforeName, c.beforeUpdate) // db.Callback().Update().After(gormCallbackName).Register(afterName, c.afterUpdate) // case "delete": // db.Callback().Delete().Before(gormCallbackName).Register(beforeName, c.beforeDelete) // db.Callback().Delete().After(gormCallbackName).Register(afterName, c.afterDelete) // case "row_query": // db.Callback().Raw().Before(gormCallbackName).Register(beforeName, c.beforeRowQuery) // db.Callback().Raw().After(gormCallbackName).Register(afterName, c.afterRowQuery) // } //}