nl-video-api/utility/database/hook.go

package database

import (
	"context"
	"database/sql/driver"
	"time"

	"github.com/gogf/gf/v2/database/gdb"
	"github.com/gogf/gf/v2/frame/g"
	"nl-video-api/utility/logger"
)

// SQLHook SQL执行钩子，用于记录SQL日志和错误
type SQLHook struct{}

// BeforeQuery 查询前钩子
func (h *SQLHook) BeforeQuery(ctx context.Context, link gdb.Link, sql string, args []interface{}) (context.Context, error) {
	// 记录查询开始时间
	ctx = context.WithValue(ctx, "sql_start_time", time.Now())
	return ctx, nil
}

// AfterQuery 查询后钩子
func (h *SQLHook) AfterQuery(ctx context.Context, link gdb.Link, sql string, args []interface{}, result gdb.Result, err error) error {
	// 计算执行时间
	startTime, ok := ctx.Value("sql_start_time").(time.Time)
	var duration time.Duration
	if ok {
		duration = time.Since(startTime)
	}

	// 获取影响行数
	var rowsAffected int64
	if result != nil {
		rowsAffected = int64(result.Len())
	}

	if err != nil {
		// 记录SQL执行错误
		logger.LogDatabaseError(ctx, "Query", sql, err)
		logger.LogError(ctx, "SQL查询失败 | SQL: %s | Args: %v | Duration: %v | Error: %v", 
			sql, args, duration, err)
	} else {
		// 记录SQL执行成功
		logger.LogDatabaseQuery(ctx, sql, args, duration, rowsAffected)
	}

	return err
}

// BeforeExec 执行前钩子
func (h *SQLHook) BeforeExec(ctx context.Context, link gdb.Link, sql string, args []interface{}) (context.Context, error) {
	// 记录执行开始时间
	ctx = context.WithValue(ctx, "sql_start_time", time.Now())
	return ctx, nil
}

// AfterExec 执行后钩子
func (h *SQLHook) AfterExec(ctx context.Context, link gdb.Link, sql string, args []interface{}, result driver.Result, err error) error {
	// 计算执行时间
	startTime, ok := ctx.Value("sql_start_time").(time.Time)
	var duration time.Duration
	if ok {
		duration = time.Since(startTime)
	}

	// 获取影响行数
	var rowsAffected int64
	if result != nil {
		if affected, e := result.RowsAffected(); e == nil {
			rowsAffected = affected
		}
	}

	if err != nil {
		// 记录SQL执行错误
		logger.LogDatabaseError(ctx, "Exec", sql, err)
		logger.LogError(ctx, "SQL执行失败 | SQL: %s | Args: %v | Duration: %v | Error: %v", 
			sql, args, duration, err)
	} else {
		// 记录SQL执行成功
		logger.LogDatabaseQuery(ctx, sql, args, duration, rowsAffected)
	}

	return err
}

// BeforePrepare 预处理前钩子
func (h *SQLHook) BeforePrepare(ctx context.Context, link gdb.Link, sql string) (context.Context, error) {
	return ctx, nil
}

// AfterPrepare 预处理后钩子
func (h *SQLHook) AfterPrepare(ctx context.Context, link gdb.Link, sql string, stmt *gdb.Stmt, err error) error {
	if err != nil {
		logger.LogDatabaseError(ctx, "Prepare", sql, err)
		logger.LogError(ctx, "SQL预处理失败 | SQL: %s | Error: %v", sql, err)
	}
	return err
}

// BeforeCommit 提交前钩子
func (h *SQLHook) BeforeCommit(ctx context.Context, link gdb.Link) (context.Context, error) {
	logger.LogSQL(ctx, "事务提交开始")
	return ctx, nil
}

// AfterCommit 提交后钩子
func (h *SQLHook) AfterCommit(ctx context.Context, link gdb.Link, err error) error {
	if err != nil {
		logger.LogError(ctx, "事务提交失败 | Error: %v", err)
	} else {
		logger.LogSQL(ctx, "事务提交成功")
	}
	return err
}

// BeforeRollback 回滚前钩子
func (h *SQLHook) BeforeRollback(ctx context.Context, link gdb.Link) (context.Context, error) {
	logger.LogSQL(ctx, "事务回滚开始")
	return ctx, nil
}

// AfterRollback 回滚后钩子
func (h *SQLHook) AfterRollback(ctx context.Context, link gdb.Link, err error) error {
	if err != nil {
		logger.LogError(ctx, "事务回滚失败 | Error: %v", err)
	} else {
		logger.LogSQL(ctx, "事务回滚成功")
	}
	return err
}

// InitDatabaseHook 初始化数据库钩子
func InitDatabaseHook() {
	// 获取默认数据库实例
	db := g.DB()
	
	// 添加SQL执行钩子
	db.AddHook(&SQLHook{})
	
	logger.LogInfo(context.Background(), "数据库钩子初始化完成")
}

// LogDatabaseConnection 记录数据库连接日志
func LogDatabaseConnection(ctx context.Context, config gdb.ConfigNode, err error) {
	if err != nil {
		logger.LogError(ctx, "数据库连接失败 | Host: %s | Database: %s | Error: %v", 
			config.Host, config.Name, err)
	} else {
		logger.LogInfo(ctx, "数据库连接成功 | Host: %s | Database: %s", 
			config.Host, config.Name)
	}
}

// LogDatabasePing 记录数据库ping日志
func LogDatabasePing(ctx context.Context, duration time.Duration, err error) {
	if err != nil {
		logger.LogError(ctx, "数据库ping失败 | Duration: %v | Error: %v", duration, err)
	} else {
		logger.LogInfo(ctx, "数据库ping成功 | Duration: %v", duration)
	}
}

// LogTransactionStart 记录事务开始日志
func LogTransactionStart(ctx context.Context, txId string) {
	logger.LogSQL(ctx, "事务开始 | TxID: %s", txId)
}

// LogTransactionEnd 记录事务结束日志
func LogTransactionEnd(ctx context.Context, txId string, success bool, duration time.Duration) {
	if success {
		logger.LogSQL(ctx, "事务结束(成功) | TxID: %s | Duration: %v", txId, duration)
	} else {
		logger.LogError(ctx, "事务结束(失败) | TxID: %s | Duration: %v", txId, duration)
	}
}