Golang中如何将中间件的r.Context()传递给数据库查询

在Golang中，将中间件的r.Context()传递给数据库查询的正确做法是通过上下文传递，而不是使用全局变量。以下是几种实现方式：

1. 通过函数参数传递（推荐）

这是最直接和类型安全的方式：

// 仓库层接口
type Repository interface {
    Add(ctx context.Context, data ValidatedData) error
}

// 仓库实现
type UserRepository struct {
    db *sql.DB
}

func (r *UserRepository) Add(ctx context.Context, data ValidatedData) error {
    query := `INSERT INTO users (name, email) VALUES ($1, $2) RETURNING id`
    
    // 使用传入的上下文
    err := r.db.QueryRowContext(ctx, query, data.Name, data.Email).Scan(&data.ID)
    if err != nil {
        return fmt.Errorf("insert user: %w", err)
    }
    return nil
}

// 处理器
func (a *Handler) AddHandler(w http.ResponseWriter, r *http.Request) {
    // 从 r 获取数据
    // 验证数据
    
    // 传递上下文到仓库
    err := a.repo.Add(r.Context(), validatedData)
    if err != nil {
        http.Error(w, err.Error(), http.StatusInternalServerError)
        return
    }
    
    w.WriteHeader(http.StatusCreated)
}

2. 使用上下文存储中间件数据

如果需要在上下文中存储额外数据：

// 定义上下文键类型
type contextKey string

const (
    userKey contextKey = "user"
    dbKey   contextKey = "db"
)

// 中间件设置数据
func AuthMiddleware(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        // 模拟获取用户信息
        user := map[string]interface{}{
            "id":    123,
            "name":  "john",
            "email": "john@example.com",
        }
        
        // 将数据存入上下文
        ctx := context.WithValue(r.Context(), userKey, user)
        next.ServeHTTP(w, r.WithContext(ctx))
    })
}

// 仓库中获取上下文数据
func (r *UserRepository) AddWithUser(ctx context.Context, data ValidatedData) error {
    // 从上下文中获取用户信息
    if user, ok := ctx.Value(userKey).(map[string]interface{}); ok {
        userID := user["id"].(int)
        query := `INSERT INTO users (name, email, created_by) VALUES ($1, $2, $3)`
        _, err := r.db.ExecContext(ctx, query, data.Name, data.Email, userID)
        return err
    }
    
    query := `INSERT INTO users (name, email) VALUES ($1, $2)`
    _, err := r.db.ExecContext(ctx, query, data.Name, data.Email)
    return err
}

3. 避免全局变量的原因

全局变量会导致：

• 并发安全问题
• 测试困难
• 代码耦合度高
• 难以追踪数据流

4. 完整示例

// main.go
func main() {
    db := setupDatabase()
    repo := NewUserRepository(db)
    handler := NewHandler(repo)
    
    mux := http.NewServeMux()
    mux.HandleFunc("/add", handler.AddHandler)
    
    // 应用中间件
    wrappedMux := ContextTimeOut(mux)
    
    http.ListenAndServe(":8080", wrappedMux)
}

// handler.go
type Handler struct {
    repo Repository
}

func NewHandler(repo Repository) *Handler {
    return &Handler{repo: repo}
}

func (h *Handler) AddHandler(w http.ResponseWriter, r *http.Request) {
    // 解析请求数据
    var data ValidatedData
    if err := json.NewDecoder(r.Body).Decode(&data); err != nil {
        http.Error(w, err.Error(), http.StatusBadRequest)
        return
    }
    
    // 验证数据
    if err := validateData(data); err != nil {
        http.Error(w, err.Error(), http.StatusBadRequest)
        return
    }
    
    // 使用上下文调用仓库
    if err := h.repo.Add(r.Context(), data); err != nil {
        // 检查是否为超时错误
        if errors.Is(err, context.DeadlineExceeded) {
            http.Error(w, "request timeout", http.StatusGatewayTimeout)
            return
        }
        http.Error(w, err.Error(), http.StatusInternalServerError)
        return
    }
    
    w.WriteHeader(http.StatusCreated)
    json.NewEncoder(w).Encode(map[string]interface{}{
        "status": "success",
        "data":   data,
    })
}

// repository.go
type UserRepository struct {
    db *sql.DB
}

func NewUserRepository(db *sql.DB) *UserRepository {
    return &UserRepository{db: db}
}

func (r *UserRepository) Add(ctx context.Context, data ValidatedData) error {
    // 使用传入的上下文执行查询
    query := `
        INSERT INTO users (name, email, created_at) 
        VALUES ($1, $2, $3) 
        RETURNING id, created_at
    `
    
    row := r.db.QueryRowContext(ctx, query, 
        data.Name, 
        data.Email, 
        time.Now(),
    )
    
    return row.Scan(&data.ID, &data.CreatedAt)
}

通过函数参数传递上下文是最佳实践，它保持了代码的清晰性、可测试性和类型安全性。中间件设置的超时上下文会自动传播到所有使用该上下文的数据库操作中。