使用Gorm结合PostgreSQL行级安全，如何为每个请求设置会话变量

我们需要找到Go/Gorm中等效于以下Java代码的实现：该Java代码会在每次从数据源获取连接时，设置PostgreSQL会话变量以实现多租户数据隔离。

// Every time the app asks the data source for a connection
// set the PostgreSQL session variable to the current tenant
// to enforce data isolation.
@Override
public Connection getConnection() throws SQLException {
    Connection connection = super.getConnection();
    try (Statement sql = connection.createStatement()) {
        sql.execute("SET app.current_tenant = '" + TenantContext.getTenant() + "'");
    }
    return connection;
}

需求说明 #

每个请求需要设置不同的app.current_tenant值，但所有请求共用同一个数据库DSN。

现有基础Gorm连接代码 #

import (      
 "gorm.io/driver/postgres"
 "gorm.io/gorm"
)

dsn := "host=localhost user=gorm password=gorm dbname=gorm port=9920 sslmode=disable TimeZone=Asia/Shanghai"
db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{})

遇到的疑问 #

• 是否需要为每个请求调用gorm.Open()创建新Session？
• 无法区分Gorm中的连接（依赖DSN）与会话（基于同一连接），因为gorm.Open()和.Session()返回相同类型
• 推测的流程是否正确：仅初始化一次全局*DB实例，每个请求通过.Session()创建会话并执行设置语句？
• 会话是否需要手动关闭？如何关闭不影响原连接池？

解决方案 #

核心结论

1. 仅初始化一次全局连接池：通过gorm.Open()创建的*gorm.DB是连接池的管理者，只需在程序启动时初始化一次，不要每次请求都重新打开。
2. 每个请求创建独立会话：通过.Session()为每个请求创建独立会话，在会话上执行SET语句，会话会绑定连接池中的一个连接，SET的变量仅在当前连接的会话周期内有效。

具体实现

1. 初始化全局数据库连接池

package main

import (
    "gorm.io/driver/postgres"
    "gorm.io/gorm"
)

var globalDB *gorm.DB

func initDB() error {
    dsn := "host=localhost user=gorm password=gorm dbname=gorm port=9920 sslmode=disable TimeZone=Asia/Shanghai"
    var err error
    globalDB, err = gorm.Open(postgres.Open(dsn), &gorm.Config{})
    if err != nil {
        return err
    }

    // 配置连接池参数（可选）
    sqlDB, err := globalDB.DB()
    if err != nil {
        return err
    }
    sqlDB.SetMaxOpenConns(100) // 最大打开连接数
    sqlDB.SetMaxIdleConns(20)  // 最大空闲连接数
    return nil
}

2. 每个请求的处理流程

import "fmt"

type User struct {
    ID   uint
    Name string
}

func handleRequest(tenantID string) error {
    // 创建独立会话，NewDB: true确保会话完全独立，不继承全局DB的上下文
    sessionDB := globalDB.Session(&gorm.Session{NewDB: true})

    // 执行SET语句：使用参数化绑定避免SQL注入，绝对不要用fmt.Sprintf拼接！
    if err := sessionDB.Exec("SET myapp.current_tenant_id = ?", tenantID).Error; err != nil {
        return fmt.Errorf("failed to set tenant: %w", err)
    }

    // 后续所有数据库操作都使用这个sessionDB实例
    var users []User
    if err := sessionDB.Find(&users).Error; err != nil {
        return fmt.Errorf("failed to query users: %w", err)
    }

    // 会话无需手动关闭：请求结束后，连接会自动归还到连接池
    return nil
}

关键要点解释

• 会话与连接的关系：
  • globalDB是连接池管理者，维护多个数据库连接；
  • 每次调用.Session()会从连接池获取一个连接，绑定到新的*gorm.DB实例（会话）；
  • SET语句的作用域是当前连接，连接归还到池后若被复用，后续请求需要重新执行SET（必须每个请求都设置，避免租户数据泄露）。
• 为什么用NewDB: true：创建完全独立的会话，确保不会继承全局DB的任何上下文（比如之前的租户变量），保证每个请求的会话干净。
• 会话关闭问题：普通会话无需手动关闭，Gorm会自动将连接归还到连接池；如果是事务会话（tx.Begin()），则需要手动调用Commit()或Rollback()。

优化：用HTTP中间件自动设置租户

可以通过中间件统一处理租户变量设置，避免重复代码：

import (
    "context"
    "net/http"
)

const dbSessionKey = "db_session"

func TenantMiddleware(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        // 从请求头/URL参数等位置获取租户ID
        tenantID := r.Header.Get("X-Tenant-ID")
        if tenantID == "" {
            http.Error(w, "tenant ID required", http.StatusBadRequest)
            return
        }

        // 创建会话并设置租户变量
        sessionDB := globalDB.Session(&gorm.Session{NewDB: true})
        if err := sessionDB.Exec("SET myapp.current_tenant_id = ?", tenantID).Error; err != nil {
            http.Error(w, "failed to initialize tenant session", http.StatusInternalServerError)
            return
        }

        // 将会话存入请求上下文，供后续处理函数使用
        ctx := context.WithValue(r.Context(), dbSessionKey, sessionDB)
        next.ServeHTTP(w, r.WithContext(ctx))
    })
}

// 示例处理函数
func getUserHandler(w http.ResponseWriter, r *http.Request) {
    sessionDB, ok := r.Context().Value(dbSessionKey).(*gorm.DB)
    if !ok {
        http.Error(w, "database session not found", http.StatusInternalServerError)
        return
    }

    var user User
    userID := r.URL.Query().Get("user_id")
    if err := sessionDB.First(&user, userID).Error; err != nil {
        http.Error(w, err.Error(), http.StatusInternalServerError)
        return
    }

    // 返回用户数据...
}

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