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

单会话模式

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

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

持续会话模式

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

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

在 Hooks/Callbacks 中使用 Context

您可以从当前 Statement中访问 Context 对象,例如︰

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

Chi 中间件示例

在处理 API 请求时持续会话模式会比较有用。例如,您可以在中间件中为 *gorm.DB 设置超时 Context,然后使用 *gorm.DB 处理所有请求

下面是一个 Chi 中间件的示例:

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

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

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

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

18. // 你的其他 DB 操作...
19. })

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

24. var user User
25. db.First(&user)

27. // 你的其他 DB 操作...
28. })

注意 通过 WithContext 设置的 Context 是线程安全的,参考会话获取详情

Logger

Logger 也可以支持 Context,可用于日志追踪,查看 Logger 获取详情