Go 的 Context 是请求级别的超时、取消、传值机制。GORM 完整支持 Context,让数据库操作可以响应请求取消。
传入 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)
}
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"})
}
创建带 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 超时会影响连接池:
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,设置合理超时,提升系统健壮性。