Rust+Axum+SeaORM多租户PostgreSQL会话保持问题求助

问题出在SeaORM默认使用连接池，每次数据库操作可能从池中获取不同的连接——你单独给某个连接设置app.tenant后，后续查询如果拿到其他连接，自然读取不到对应租户的数据。要解决这个问题，核心是确保同一个请求的所有数据库操作都在同一个连接上执行，并且在该连接上绑定租户变量。

以下是几种可行的方案：

方案一：使用事务 + SET LOCAL（推荐） #

SET LOCAL是PostgreSQL的会话变量设置命令，仅在当前事务内有效。事务结束后，连接放回池时该设置会自动失效，不会污染其他请求的连接，安全性更高。

代码示例： #

use axum::{extract::Path, Json, State};
use sea_orm::{DatabaseConnection, DbErr, EntityTrait};
use std::sync::Arc;

// 假设你的Ticket实体定义
use crate::entities::ticket;

async fn get_tickets(
    Path(tenant_id): Path<i32>,
    State(db): State<Arc<DatabaseConnection>>,
) -> Result<Json<Vec<ticket::Model>>, DbErr> {
    let tickets = db.transaction(|tx| async move {
        // 在事务内设置租户变量
        sea_orm::sqlx::query("SET LOCAL app.tenant = $1")
            .bind(tenant_id)
            .execute(tx)
            .await?;
        
        // 在同一个事务中查询数据
        ticket::Entity::find().all(tx).await
    }).await?;

    Ok(Json(tickets))
}

方案二：获取单个连接并复用 #

直接从连接池中获取一个单独的连接，在该连接上设置租户变量，然后所有查询都使用这个连接。注意用完后建议重置租户变量，避免连接放回池后污染后续请求。

代码示例： #

async fn get_tickets(
    Path(tenant_id): Path<i32>,
    State(db): State<Arc<DatabaseConnection>>,
) -> Result<Json<Vec<ticket::Model>>, DbErr> {
    // 获取单个连接
    let mut conn = db.get_connection().await?;

    // 设置租户变量
    sea_orm::sqlx::query("SET app.tenant = $1")
        .bind(tenant_id)
        .execute(&mut conn)
        .await?;

    // 使用同一个连接执行查询
    let tickets = ticket::Entity::find().all(&mut conn).await?;

    // 重置租户变量，避免连接污染
    sea_orm::sqlx::query("SET app.tenant = NULL")
        .execute(&mut conn)
        .await?;

    Ok(Json(tickets))
}

方案三：自定义连接提取器 #

如果多个请求都需要处理租户，可以封装一个自定义Axum提取器，自动完成租户提取和连接配置：

use axum::{extract::FromRequestParts, http::request::Parts, State};
use sea_orm::{DatabaseConnection, DbErr};
use std::sync::Arc;
use async_trait::async_trait;

#[derive(Clone)]
pub struct TenantDb {
    pub conn: Arc<DatabaseConnection>,
    pub tenant_id: i32,
}

#[async_trait]
impl<S> FromRequestParts<S> for TenantDb
where
    S: Send + Sync,
{
    type Rejection = DbErr;

    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
        // 示例从请求扩展数据中获取租户ID，实际可替换为从请求头/JWT解析
        let tenant_id = parts
            .extensions
            .get::<i32>()
            .ok_or(DbErr::Custom("Tenant ID not found".into()))?
            .clone();
        
        let conn = state
            .get::<Arc<DatabaseConnection>>()
            .ok_or(DbErr::Custom("Database connection not in state".into()))?
            .clone();

        Ok(Self { conn, tenant_id })
    }
}

// 使用提取器的处理函数
async fn get_tickets(tenant_db: TenantDb) -> Result<Json<Vec<ticket::Model>>, DbErr> {
    let tickets = tenant_db.conn.transaction(|tx| async move {
        sea_orm::sqlx::query("SET LOCAL app.tenant = $1")
            .bind(tenant_db.tenant_id)
            .execute(tx)
            .await?;
        
        ticket::Entity::find().all(tx).await
    }).await?;

    Ok(Json(tickets))
}

关键注意点 #

• 避免直接通过&db执行查询，因为每次都会从池里取新连接，之前设置的会话变量不会生效。
• 优先使用SET LOCAL + 事务的方案，既保证会话变量仅在当前请求有效，又避免连接池污染。
• 租户ID的提取要做安全校验，比如通过JWT签名验证，防止恶意篡改。

内容的提问来源于stack exchange，提问作者Luciano Vieira Rodrigues