Context

GORM 通过 WithContext 方法提供了 Context 支持

单会话模式

单会话模式通常被用于执行单次操作

db.WithContext(ctx).Find(&users)

持续会话模式

持续会话模式通常被用于执行一系列操作，例如：

tx := db.WithContext(ctx)
tx.First(&user, 1)
tx.Model(&user).Update("role", "admin")

Context timeout

You could passing in a context with a timeout to db.WithContext to set timeout for long running queries, for example:

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

db.WithContext(ctx).Find(&users)

Context in Hooks/Callbacks

You could access the Context object from current Statement, for example:

func (u *User) BeforeCreate(tx *gorm.DB) (err error) {
  ctx := tx.Statement.Context
  // ...
  return
}

Chi Middleware Example

Continuous session mode which might be helpful when handling API requests, for example, you can set up *gorm.DB with Timeout Context in middlewares, and then use the *gorm.DB when processing all requests

Following is a Chi middleware example:

func SetDBMiddleware(next http.Handler) http.Handler {
  return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    timeoutContext, _ := context.WithTimeout(context.Background(), time.Second)
    ctx := context.WithValue(r.Context(), "DB", db.WithContext(timeoutContext))
    next.ServeHTTP(w, r.WithContext(ctx))
  })
}

r := chi.NewRouter()
r.Use(SetDBMiddleware)

r.Get("/", func(w http.ResponseWriter, r *http.Request) {
  db, ok := ctx.Value("DB").(*gorm.DB)

  var users []User
  db.Find(&users)

  // lots of db operations
})

r.Get("/user", func(w http.ResponseWriter, r *http.Request) {
  db, ok := ctx.Value("DB").(*gorm.DB)

  var user User
  db.First(&user)

  // lots of db operations
})

NOTE Set Context with WithContext is goroutine-safe, refer Session for details

Logger

Logger accepts Context too, you can use it for log tracking, refer Logger for details

GORM