上下文支持

Go 的 Context 是请求级别的超时、取消、传值机制。GORM 完整支持 Context,让数据库操作可以响应请求取消。

WithContext

传入 Context:

ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()

var users []User
err := db.WithContext(ctx).Find(&users).Error

超时后操作会被取消,返回 context 错误。

请求超时

HTTP 请求中设置超时:

func GetUser(c *gin.Context) {
    ctx, cancel := context.WithTimeout(c.Request.Context(), 3*time.Second)
    defer cancel()
    
    var user User
    if err := db.WithContext(ctx).First(&user, c.Param("id")).Error; err != nil {
        if errors.Is(err, context.DeadlineExceeded) {
            c.JSON(408, gin.H{"error": "请求超时"})
            return
        }
        c.JSON(500, gin.H{"error": err.Error()})
        return
    }
    
    c.JSON(200, user)
}

请求取消

用户取消请求时自动中断数据库操作:

func SearchUsers(c *gin.Context) {
    ctx := c.Request.Context()
    
    var users []User
    if err := db.WithContext(ctx).Where("name LIKE ?", c.Query("q")).Find(&users).Error; err != nil {
        if errors.Is(err, context.Canceled) {
            return // 客户端已断开,直接返回
        }
        c.JSON(500, gin.H{"error": err.Error()})
        return
    }
    
    c.JSON(200, users)
}

传递值

Context 可以传递请求级别的数据:

type contextKey string

const userIDKey contextKey = "user_id"

func AuthMiddleware() gin.HandlerFunc {
    return func(c *gin.Context) {
        userID := parseToken(c.GetHeader("Authorization"))
        ctx := context.WithValue(c.Request.Context(), userIDKey, userID)
        c.Request = c.Request.WithContext(ctx)
        c.Next()
    }
}

func (u *User) BeforeCreate(tx *gorm.DB) error {
    ctx := tx.Statement.Context
    if userID, ok := ctx.Value(userIDKey).(uint); ok {
        u.CreatedBy = userID
    }
    return nil
}

func CreateUser(c *gin.Context) {
    user := User{Name: c.PostForm("name")}
    db.WithContext(c.Request.Context()).Create(&user)
}

事务中的 Context

func CreateOrder(c *gin.Context) {
    ctx := c.Request.Context()
    
    err := db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
        var user User
        if err := tx.First(&user, c.GetUint("user_id")).Error; err != nil {
            return err
        }
        
        order := Order{UserID: user.ID}
        if err := tx.Create(&order).Error; err != nil {
            return err
        }
        
        return nil
    })
    
    if err != nil {
        c.JSON(500, gin.H{"error": err.Error()})
        return
    }
    
    c.JSON(200, gin.H{"message": "success"})
}

Session 复用

创建带 Context 的 Session:

func NewDBWithContext(ctx context.Context, db *gorm.DB) *gorm.DB {
    return db.WithContext(ctx)
}

func ProcessRequest(ctx context.Context, db *gorm.DB) error {
    db = db.WithContext(ctx)
    
    var users []User
    if err := db.Find(&users).Error; err != nil {
        return err
    }
    
    for _, user := range users {
        if err := db.Model(&user).Update("processed", true).Error; err != nil {
            return err
        }
    }
    
    return nil
}

超时策略

不同操作设置不同超时:

func QueryWithTimeout(db *gorm.DB, timeout time.Duration, fn func(*gorm.DB) error) error {
    ctx, cancel := context.WithTimeout(context.Background(), timeout)
    defer cancel()
    return fn(db.WithContext(ctx))
}

// 快速查询
QueryWithTimeout(db, 2*time.Second, func(db *gorm.DB) error {
    return db.First(&user, 1).Error
})

// 慢查询
QueryWithTimeout(db, 30*time.Second, func(db *gorm.DB) error {
    return db.Where("complex condition").Find(&results).Error
})

Context 与连接池

Context 超时会影响连接池:

ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()

// 如果连接池满了,等待连接会受超时限制
db.WithContext(ctx).Find(&users)

实际案例

API 请求处理

func GetUserHandler(c *gin.Context) {
    ctx, cancel := context.WithTimeout(c.Request.Context(), 5*time.Second)
    defer cancel()
    
    id := c.Param("id")
    
    var user User
    if err := db.WithContext(ctx).Preload("Orders").First(&user, id).Error; err != nil {
        if errors.Is(err, context.DeadlineExceeded) {
            c.JSON(408, gin.H{"error": "请求超时"})
            return
        }
        if errors.Is(err, gorm.ErrRecordNotFound) {
            c.JSON(404, gin.H{"error": "用户不存在"})
            return
        }
        c.JSON(500, gin.H{"error": "服务器错误"})
        return
    }
    
    c.JSON(200, user)
}

后台任务

func BackgroundTask(db *gorm.DB) {
    ctx := context.Background()
    
    for {
        select {
        case <-time.After(10 * time.Second):
            var pending []Order
            db.WithContext(ctx).Where("status = ?", "pending").Find(&pending)
            // 处理订单...
        }
    }
}

func GracefulShutdown(db *gorm.DB, shutdownCtx context.Context) {
    <-shutdownCtx.Done()
    
    ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
    defer cancel()
    
    sqlDB, _ := db.WithContext(ctx).DB()
    sqlDB.Close()
}

小结

Context 支持让 GORM 操作可以响应超时和取消。在 HTTP 请求中传递 Context,设置合理超时,提升系统健壮性。