如何向AuthMw中间件传递数据库实例以实现数据库操作

我来帮你搞定这个问题，在Go里把main函数里初始化的数据库实例传递给中间件，有几种实用的方法，我给你一一拆解：

方法1：闭包注入（最推荐） #

这是最直观也最符合依赖注入思想的方式，我们可以把AuthMw改造成一个接收数据库实例作为参数的函数，让它返回一个标准的中间件函数。这样中间件就能直接持有db的引用，类型安全还容易测试。

改造后的代码示例：

import (
    "fmt"
    "net/http"
    "github.com/gorilla/mux"
    "gorm.io/gorm"
)

func main() {
    // 初始化数据库
    db, err := gorm.Open("postgres", "your-dsn-here")
    if err != nil {
        // 记得处理初始化错误，比如直接终止程序
        panic(fmt.Sprintf("failed to connect database: %v", err))
    }

    r := mux.NewRouter()
    r.Handle("/ws", serveWebsocket(hub))
    // 把db传给AuthMw，得到一个可用的中间件
    r.Use(AuthMw(db))
    // ... 其他路由配置
}

// 改造AuthMw：接收db参数，返回标准中间件函数
func AuthMw(db *gorm.DB) func(next http.Handler) http.Handler {
    return func(next http.Handler) http.Handler {
        return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
            token := r.URL.Query().Get("token")
            fmt.Printf("AuthMiddleware token is: %v\n", token)

            // 现在可以直接在这里使用db做数据库操作了
            // 示例：根据token查询用户
            type User struct {
                ID    uint
                Token string
            }
            var user User
            result := db.Where("token = ?", token).First(&user)
            if result.Error != nil {
                // 认证失败，返回401
                http.Error(w, "Unauthorized", http.StatusUnauthorized)
                return
            }

            // 认证通过，继续执行后续处理逻辑
            next.ServeHTTP(w, r)
        })
    }
}

方法2：通过Context传递 #

如果你没办法修改中间件的签名（比如中间件是第三方库提供的），可以把db存入请求的Context中，在中间件里再取出来。不过这种方法需要类型断言，类型安全性稍差，适合临时场景。

示例代码：

import (
    "context"
    "fmt"
    "net/http"
    "github.com/gorilla/mux"
    "gorm.io/gorm"
)

// 最好定义一个自定义的context key，避免和其他库冲突
type dbContextKey struct{}

func main() {
    db, err := gorm.Open("postgres", "your-dsn-here")
    if err != nil {
        panic(fmt.Sprintf("failed to connect database: %v", err))
    }

    r := mux.NewRouter()
    r.Handle("/ws", serveWebsocket(hub))
    
    // 先添加一个中间件，把db存入请求的Context
    r.Use(func(next http.Handler) http.Handler {
        return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
            ctx := context.WithValue(r.Context(), dbContextKey{}, db)
            next.ServeHTTP(w, r.WithContext(ctx))
        })
    })
    
    // 再使用原来的AuthMw
    r.Use(AuthMw)
    // ...
}

func AuthMw(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        token := r.URL.Query().Get("token")
        fmt.Printf("AuthMiddleware token is: %v\n", token)

        // 从Context中取出db，需要类型断言
        db, ok := r.Context().Value(dbContextKey{}).(*gorm.DB)
        if !ok {
            // 取出失败，返回500错误
            http.Error(w, "Internal Server Error", http.StatusInternalServerError)
            return
        }

        // 接下来就可以用db做操作了
        // ...

        next.ServeHTTP(w, r)
    })
}

方法3：全局变量（不推荐） #

最简单但最不推荐的方式是把db定义为全局变量，这样中间件可以直接访问。但全局变量会导致代码耦合度高，测试时很难替换mock实例，大型项目里尽量避免。

示例代码：

import (
    "fmt"
    "net/http"
    "github.com/gorilla/mux"
    "gorm.io/gorm"
)

// 全局数据库变量
var db *gorm.DB

func main() {
    var err error
    db, err = gorm.Open("postgres", "your-dsn-here")
    if err != nil {
        panic(fmt.Sprintf("failed to connect database: %v", err))
    }

    r := mux.NewRouter()
    r.Handle("/ws", serveWebsocket(hub))
    r.Use(AuthMw)
    // ...
}

func AuthMw(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        token := r.URL.Query().Get("token")
        fmt.Printf("AuthMiddleware token is: %v\n", token)

        // 直接使用全局db变量
        // ...

        next.ServeHTTP(w, r)
    })
}

总结 #

优先选择方法1的闭包注入，它既保证了类型安全，又让依赖关系清晰，后续维护和测试都更方便。

内容的提问来源于stack exchange，提问作者Blankman