gopkgsv2/trace/gorm.go

// 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)
//	}
//}