321 lines
8.5 KiB
Go
321 lines
8.5 KiB
Go
|
package tracing
|
|||
|
|
|
|||
|
|
import (
|
|||
|
|
"context"
|
|||
|
|
"fmt"
|
|||
|
|
"strings"
|
|||
|
|
"time"
|
|||
|
|
|
|||
|
|
"go.opentelemetry.io/otel/attribute"
|
|||
|
|
"go.opentelemetry.io/otel/codes"
|
|||
|
|
"go.opentelemetry.io/otel/trace"
|
|||
|
|
"go.uber.org/zap"
|
|||
|
|
"gorm.io/gorm"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
const (
|
|||
|
|
gormSpanKey = "otel:span"
|
|||
|
|
gormOperationKey = "otel:operation"
|
|||
|
|
gormTableNameKey = "otel:table_name"
|
|||
|
|
gormStartTimeKey = "otel:start_time"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
// GormTracingPlugin GORM链路追踪插件
|
|||
|
|
type GormTracingPlugin struct {
|
|||
|
|
tracer *Tracer
|
|||
|
|
logger *zap.Logger
|
|||
|
|
config GormPluginConfig
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// GormPluginConfig GORM插件配置
|
|||
|
|
type GormPluginConfig struct {
|
|||
|
|
IncludeSQL bool
|
|||
|
|
IncludeValues bool
|
|||
|
|
SlowThreshold time.Duration
|
|||
|
|
ExcludeTables []string
|
|||
|
|
SanitizeSQL bool
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// DefaultGormPluginConfig 默认GORM插件配置
|
|||
|
|
func DefaultGormPluginConfig() GormPluginConfig {
|
|||
|
|
return GormPluginConfig{
|
|||
|
|
IncludeSQL: true,
|
|||
|
|
IncludeValues: false, // 生产环境建议设为false避免记录敏感数据
|
|||
|
|
SlowThreshold: 200 * time.Millisecond,
|
|||
|
|
ExcludeTables: []string{"migrations", "schema_migrations"},
|
|||
|
|
SanitizeSQL: true,
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// NewGormTracingPlugin 创建GORM追踪插件
|
|||
|
|
func NewGormTracingPlugin(tracer *Tracer, logger *zap.Logger) *GormTracingPlugin {
|
|||
|
|
return &GormTracingPlugin{
|
|||
|
|
tracer: tracer,
|
|||
|
|
logger: logger,
|
|||
|
|
config: DefaultGormPluginConfig(),
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// Name 返回插件名称
|
|||
|
|
func (p *GormTracingPlugin) Name() string {
|
|||
|
|
return "gorm-otel-tracing"
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// Initialize 初始化插件
|
|||
|
|
func (p *GormTracingPlugin) Initialize(db *gorm.DB) error {
|
|||
|
|
// 注册各种操作的回调
|
|||
|
|
callbacks := []string{"create", "query", "update", "delete", "raw"}
|
|||
|
|
|
|||
|
|
for _, operation := range callbacks {
|
|||
|
|
switch operation {
|
|||
|
|
case "create":
|
|||
|
|
err := db.Callback().Create().Before("gorm:create").
|
|||
|
|
Register(p.Name()+":before_create", p.beforeOperation)
|
|||
|
|
if err != nil {
|
|||
|
|
return fmt.Errorf("failed to register before create callback: %w", err)
|
|||
|
|
}
|
|||
|
|
err = db.Callback().Create().After("gorm:create").
|
|||
|
|
Register(p.Name()+":after_create", p.afterOperation)
|
|||
|
|
if err != nil {
|
|||
|
|
return fmt.Errorf("failed to register after create callback: %w", err)
|
|||
|
|
}
|
|||
|
|
case "query":
|
|||
|
|
err := db.Callback().Query().Before("gorm:query").
|
|||
|
|
Register(p.Name()+":before_query", p.beforeOperation)
|
|||
|
|
if err != nil {
|
|||
|
|
return fmt.Errorf("failed to register before query callback: %w", err)
|
|||
|
|
}
|
|||
|
|
err = db.Callback().Query().After("gorm:query").
|
|||
|
|
Register(p.Name()+":after_query", p.afterOperation)
|
|||
|
|
if err != nil {
|
|||
|
|
return fmt.Errorf("failed to register after query callback: %w", err)
|
|||
|
|
}
|
|||
|
|
case "update":
|
|||
|
|
err := db.Callback().Update().Before("gorm:update").
|
|||
|
|
Register(p.Name()+":before_update", p.beforeOperation)
|
|||
|
|
if err != nil {
|
|||
|
|
return fmt.Errorf("failed to register before update callback: %w", err)
|
|||
|
|
}
|
|||
|
|
err = db.Callback().Update().After("gorm:update").
|
|||
|
|
Register(p.Name()+":after_update", p.afterOperation)
|
|||
|
|
if err != nil {
|
|||
|
|
return fmt.Errorf("failed to register after update callback: %w", err)
|
|||
|
|
}
|
|||
|
|
case "delete":
|
|||
|
|
err := db.Callback().Delete().Before("gorm:delete").
|
|||
|
|
Register(p.Name()+":before_delete", p.beforeOperation)
|
|||
|
|
if err != nil {
|
|||
|
|
return fmt.Errorf("failed to register before delete callback: %w", err)
|
|||
|
|
}
|
|||
|
|
err = db.Callback().Delete().After("gorm:delete").
|
|||
|
|
Register(p.Name()+":after_delete", p.afterOperation)
|
|||
|
|
if err != nil {
|
|||
|
|
return fmt.Errorf("failed to register after delete callback: %w", err)
|
|||
|
|
}
|
|||
|
|
case "raw":
|
|||
|
|
err := db.Callback().Raw().Before("gorm:raw").
|
|||
|
|
Register(p.Name()+":before_raw", p.beforeOperation)
|
|||
|
|
if err != nil {
|
|||
|
|
return fmt.Errorf("failed to register before raw callback: %w", err)
|
|||
|
|
}
|
|||
|
|
err = db.Callback().Raw().After("gorm:raw").
|
|||
|
|
Register(p.Name()+":after_raw", p.afterOperation)
|
|||
|
|
if err != nil {
|
|||
|
|
return fmt.Errorf("failed to register after raw callback: %w", err)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
p.logger.Info("GORM追踪插件已初始化")
|
|||
|
|
return nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// beforeOperation 操作前回调
|
|||
|
|
func (p *GormTracingPlugin) beforeOperation(db *gorm.DB) {
|
|||
|
|
// 检查是否应该跳过追踪
|
|||
|
|
if p.shouldSkipTracing(db) {
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
ctx := db.Statement.Context
|
|||
|
|
if ctx == nil {
|
|||
|
|
ctx = context.Background()
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 获取操作信息
|
|||
|
|
operation := p.getOperationType(db)
|
|||
|
|
tableName := p.getTableName(db)
|
|||
|
|
|
|||
|
|
// 检查是否应该排除此表
|
|||
|
|
if p.isExcludedTable(tableName) {
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 开始追踪
|
|||
|
|
ctx, span := p.tracer.StartDBSpan(ctx, operation, tableName)
|
|||
|
|
|
|||
|
|
// 添加基础属性
|
|||
|
|
p.tracer.AddSpanAttributes(span,
|
|||
|
|
attribute.String("db.system", "postgresql"),
|
|||
|
|
attribute.String("db.operation", operation),
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if tableName != "" {
|
|||
|
|
p.tracer.AddSpanAttributes(span, attribute.String("db.table", tableName))
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 保存追踪信息到GORM context
|
|||
|
|
db.Set(gormSpanKey, span)
|
|||
|
|
db.Set(gormOperationKey, operation)
|
|||
|
|
db.Set(gormTableNameKey, tableName)
|
|||
|
|
db.Set(gormStartTimeKey, time.Now())
|
|||
|
|
|
|||
|
|
// 更新statement context
|
|||
|
|
db.Statement.Context = ctx
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// afterOperation 操作后回调
|
|||
|
|
func (p *GormTracingPlugin) afterOperation(db *gorm.DB) {
|
|||
|
|
// 获取span
|
|||
|
|
spanValue, exists := db.Get(gormSpanKey)
|
|||
|
|
if !exists {
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
span, ok := spanValue.(trace.Span)
|
|||
|
|
if !ok {
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
defer span.End()
|
|||
|
|
|
|||
|
|
// 获取操作信息
|
|||
|
|
operation, _ := db.Get(gormOperationKey)
|
|||
|
|
tableName, _ := db.Get(gormTableNameKey)
|
|||
|
|
startTime, _ := db.Get(gormStartTimeKey)
|
|||
|
|
|
|||
|
|
// 计算执行时间
|
|||
|
|
var duration time.Duration
|
|||
|
|
if st, ok := startTime.(time.Time); ok {
|
|||
|
|
duration = time.Since(st)
|
|||
|
|
p.tracer.AddSpanAttributes(span,
|
|||
|
|
attribute.Int64("db.duration_ms", duration.Milliseconds()),
|
|||
|
|
)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 添加SQL信息
|
|||
|
|
if p.config.IncludeSQL && db.Statement.SQL.String() != "" {
|
|||
|
|
sql := db.Statement.SQL.String()
|
|||
|
|
if p.config.SanitizeSQL {
|
|||
|
|
sql = p.sanitizeSQL(sql)
|
|||
|
|
}
|
|||
|
|
p.tracer.AddSpanAttributes(span, attribute.String("db.statement", sql))
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 添加影响行数
|
|||
|
|
if db.Statement.RowsAffected >= 0 {
|
|||
|
|
p.tracer.AddSpanAttributes(span,
|
|||
|
|
attribute.Int64("db.rows_affected", db.Statement.RowsAffected),
|
|||
|
|
)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 处理错误
|
|||
|
|
if db.Error != nil {
|
|||
|
|
p.tracer.SetSpanError(span, db.Error)
|
|||
|
|
span.SetStatus(codes.Error, db.Error.Error())
|
|||
|
|
|
|||
|
|
p.logger.Error("数据库操作失败",
|
|||
|
|
zap.String("operation", fmt.Sprintf("%v", operation)),
|
|||
|
|
zap.String("table", fmt.Sprintf("%v", tableName)),
|
|||
|
|
zap.Error(db.Error),
|
|||
|
|
zap.String("trace_id", p.tracer.GetTraceID(db.Statement.Context)),
|
|||
|
|
)
|
|||
|
|
} else {
|
|||
|
|
p.tracer.SetSpanSuccess(span)
|
|||
|
|
span.SetStatus(codes.Ok, "success")
|
|||
|
|
|
|||
|
|
// 检查慢查询
|
|||
|
|
if duration > p.config.SlowThreshold {
|
|||
|
|
p.tracer.AddSpanAttributes(span,
|
|||
|
|
attribute.Bool("db.slow_query", true),
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
p.logger.Warn("慢SQL查询检测",
|
|||
|
|
zap.String("operation", fmt.Sprintf("%v", operation)),
|
|||
|
|
zap.String("table", fmt.Sprintf("%v", tableName)),
|
|||
|
|
zap.Duration("duration", duration),
|
|||
|
|
zap.String("sql", db.Statement.SQL.String()),
|
|||
|
|
zap.String("trace_id", p.tracer.GetTraceID(db.Statement.Context)),
|
|||
|
|
)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// shouldSkipTracing 检查是否应该跳过追踪
|
|||
|
|
func (p *GormTracingPlugin) shouldSkipTracing(db *gorm.DB) bool {
|
|||
|
|
// 检查是否已有span（避免重复追踪）
|
|||
|
|
if _, exists := db.Get(gormSpanKey); exists {
|
|||
|
|
return true
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
return false
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// getOperationType 获取操作类型
|
|||
|
|
func (p *GormTracingPlugin) getOperationType(db *gorm.DB) string {
|
|||
|
|
switch db.Statement.ReflectValue.Kind() {
|
|||
|
|
default:
|
|||
|
|
sql := strings.ToUpper(strings.TrimSpace(db.Statement.SQL.String()))
|
|||
|
|
if sql == "" {
|
|||
|
|
return "unknown"
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
if strings.HasPrefix(sql, "SELECT") {
|
|||
|
|
return "select"
|
|||
|
|
} else if strings.HasPrefix(sql, "INSERT") {
|
|||
|
|
return "insert"
|
|||
|
|
} else if strings.HasPrefix(sql, "UPDATE") {
|
|||
|
|
return "update"
|
|||
|
|
} else if strings.HasPrefix(sql, "DELETE") {
|
|||
|
|
return "delete"
|
|||
|
|
} else if strings.HasPrefix(sql, "CREATE") {
|
|||
|
|
return "create"
|
|||
|
|
} else if strings.HasPrefix(sql, "DROP") {
|
|||
|
|
return "drop"
|
|||
|
|
} else if strings.HasPrefix(sql, "ALTER") {
|
|||
|
|
return "alter"
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
return "query"
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// getTableName 获取表名
|
|||
|
|
func (p *GormTracingPlugin) getTableName(db *gorm.DB) string {
|
|||
|
|
if db.Statement.Table != "" {
|
|||
|
|
return db.Statement.Table
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
if db.Statement.Schema != nil && db.Statement.Schema.Table != "" {
|
|||
|
|
return db.Statement.Schema.Table
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
return ""
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// isExcludedTable 检查是否为排除的表
|
|||
|
|
func (p *GormTracingPlugin) isExcludedTable(tableName string) bool {
|
|||
|
|
for _, excluded := range p.config.ExcludeTables {
|
|||
|
|
if tableName == excluded {
|
|||
|
|
return true
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
return false
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// sanitizeSQL 清理SQL语句，移除敏感信息
|
|||
|
|
func (p *GormTracingPlugin) sanitizeSQL(sql string) string {
|
|||
|
|
// 简单的SQL清理，将参数替换为占位符
|
|||
|
|
// 在生产环境中，您可能需要更复杂的清理逻辑
|
|||
|
|
return strings.ReplaceAll(sql, "'", "?")
|
|||
|
|
}
|