如何在使用GORM时自定义连接对象并实现查询额外处理？

1. 实现自定义Conn包装器 #

定义一个包装底层sql.Conn的结构体，实现sql.Conn接口的所有方法，在需要拦截的查询/执行步骤中加入自定义逻辑（比如授权验证）：

import (
  "context"
  "database/sql"
  "errors"
  "gorm.io/gorm"
)

// CustomConn 自定义连接包装器，包含底层数据库连接
type CustomConn struct {
  underlyingConn sql.Conn
}

// PingContext 包装Ping方法，加入授权验证
func (c *CustomConn) PingContext(ctx context.Context) error {
  if err := c.authCheck(ctx); err != nil {
    return err
  }
  return c.underlyingConn.PingContext(ctx)
}

// ExecContext 拦截执行语句，加入前置处理
func (c *CustomConn) ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error) {
  if err := c.authCheck(ctx); err != nil {
    return nil, err
  }
  return c.underlyingConn.ExecContext(ctx, query, args...)
}

// QueryContext 拦截查询语句，加入前置处理
func (c *CustomConn) QueryContext(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error) {
  if err := c.authCheck(ctx); err != nil {
    return nil, err
  }
  return c.underlyingConn.QueryContext(ctx, query, args...)
}

// 必须完整实现sql.Conn的所有方法，以下为示例其余方法
func (c *CustomConn) PrepareContext(ctx context.Context, query string) (*sql.Stmt, error) {
  return c.underlyingConn.PrepareContext(ctx, query)
}

func (c *CustomConn) BeginTx(ctx context.Context, opts *sql.TxOptions) (*sql.Tx, error) {
  return c.underlyingConn.BeginTx(ctx, opts)
}

func (c *CustomConn) Close() error {
  return c.underlyingConn.Close()
}

// authCheck 自定义授权验证逻辑，对接第三方服务
func (c *CustomConn) authCheck(ctx context.Context) error {
  // 替换为实际的第三方授权验证逻辑
  // if !thirdPartyAuth.Validate(ctx) {
  //   return errors.New("authorization failed: no permission to execute query")
  // }
  return nil
}

2. 自定义Dialector重写连接创建逻辑 #

GORM通过Dialector创建连接，因此需要包装原有Dialector（如MySQL/PostgreSQL的Dialector），重写Conn方法返回自定义的CustomConn：

// CustomDialector 自定义Dialector，包装原有Dialector
type CustomDialector struct {
  originalDialector gorm.Dialector
}

// NewCustomDialector 初始化自定义Dialector
func NewCustomDialector(original gorm.Dialector) *CustomDialector {
  return &CustomDialector{originalDialector: original}
}

// 实现gorm.Dialector接口的必要方法，大部分直接委托给原Dialector
func (d *CustomDialector) Name() string {
  return d.originalDialector.Name()
}

func (d *CustomDialector) Initialize(db *gorm.DB) error {
  return d.originalDialector.Initialize(db)
}

func (d *CustomDialector) Migrator(db *gorm.DB) gorm.Migrator {
  return d.originalDialector.Migrator(db)
}

// 重写Conn方法，替换为自定义连接实例
func (d *CustomDialector) Conn(ctx context.Context) (gorm.Conn, error) {
  originalConn, err := d.originalDialector.Conn(ctx)
  if err != nil {
    return nil, err
  }
  
  // 将原连接转换为sql.Conn类型（GORM的Conn接口兼容sql.Conn）
  sqlConn, ok := originalConn.(sql.Conn)
  if !ok {
    return nil, gorm.ErrInvalidConn
  }
  
  return &CustomConn{underlyingConn: sqlConn}, nil
}

3. 初始化GORM时使用自定义Dialector #

用自定义Dialector替换原有Dialector，完成GORM初始化：

import "gorm.io/driver/mysql"

func main() {
  // 初始化原有数据库Dialector（以MySQL为例）
  originalDialector := mysql.Open("user:password@tcp(127.0.0.1:3306)/dbname?charset=utf8mb4&parseTime=True&loc=Local")
  
  // 创建自定义Dialector
  customDialector := NewCustomDialector(originalDialector)
  
  // 打开GORM连接
  db, err := gorm.Open(customDialector, &gorm.Config{})
  if err != nil {
    panic("failed to connect database")
  }
  
  // 后续所有DB操作都会经过CustomConn的拦截处理
  var user struct{ ID uint; Name string }
  db.First(&user, 1) // 触发QueryContext中的授权验证
}

注意事项 #

• 必须完整实现sql.Conn的所有方法，否则GORM无法正常处理连接生命周期
• 若需要拦截事务操作，可在BeginTx方法中加入自定义逻辑
• 自定义Dialector需实现gorm.Dialector的全部接口，非核心方法直接委托给原Dialector即可

内容的提问来源于stack exchange，提问作者mic